Medium

题目描述

给定一个可能含有重复元素的整数数组 nums，要求随机输出给定的目标数字 target 的索引。题目保证给定的数字一定存在于数组中。

实现 Solution 类：

  • Solution(int[] nums) 使用数组 nums 初始化对象。
  • int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i。如果存在多个有效的 i，则每个索引被返回的概率应当相等。

示例 1：

输入
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [1], [3]]
输出
[null, 4, 0, 2]

解释
Solution solution = new Solution([1, 2, 3, 3, 3]);
solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引被返回的概率应该相等。
solution.pick(1); // 返回 0。因为在数组中只有 nums[0] 等于 1。
solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引被返回的概率应该相等。

提示：

  • 1 <= nums.length <= 2 * 10^4
  • -2^31 <= nums[i] <= 2^31 - 1
  • target 是 nums 中的一个整数
  • 最多调用 pick 函数 10^4 次

解题思路

本题有两种主要解法：哈希表预处理和水塘抽样。

解法一：哈希表预处理

在构造函数中，使用哈希表存储每个数字对应的所有索引位置。查询时直接从对应的索引列表中随机选择一个返回。这种方法空间复杂度较高，但查询速度快。

解法二：水塘抽样（推荐）

这是一种经典的流式数据随机抽样算法。核心思想是：遍历数组时，对于第 k 个目标元素，以 1/k 的概率选择它作为结果。

具体算法：

  1. 初始化计数器 count = 0 和结果 result = -1
  2. 遍历数组，当遇到目标值时：
    • 计数器 count++
    • 以 1/count 的概率更新结果为当前索引
  3. 返回最终结果

数学证明：对于有 n 个目标元素的情况，每个元素被选中的概率都是 1/n，满足等概率要求。

水塘抽样的优势在于空间复杂度为 O(1)，特别适合处理大数据流场景。虽然时间复杂度稍高，但在数组规模不大的情况下是最优解法。

代码实现

class Solution {
private:
    vector<int> nums;
    
public:
    Solution(vector<int>& nums) : nums(nums) {
    }
    
    int pick(int target) {
        int count = 0;
        int result = -1;
        
        for (int i = 0; i < nums.size(); i++) {
            if (nums[i] == target) {
                count++;
                if (rand() % count == 0) {
                    result = i;
                }
            }
        }
        
        return result;
    }
};

import random

class Solution:

    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        count = 0
        result = -1
        
        for i, num in enumerate(self.nums):
            if num == target:
                count += 1
                if random.randint(1, count) == 1:
                    result = i
                    
        return result

public class Solution {
    private int[] nums;
    private Random rand;

    public Solution(int[] nums) {
        this.nums = nums;
        this.rand = new Random();
    }
    
    public int Pick(int target) {
        int count = 0;
        int result = -1;
        
        for (int i = 0; i < nums.Length; i++) {
            if (nums[i] == target) {
                count++;
                if (rand.Next(count) == 0) {
                    result = i;
                }
            }
        }
        
        return result;
    }
}

var Solution = function(nums) {
    this.nums = nums;
};

Solution.prototype.pick = function(target) {
    let count = 0;
    let result = -1;
    
    for (let i = 0; i < this.nums.length; i++) {
        if (this.nums[i]
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

方法时间复杂度空间复杂度
哈希表预处理初始化: O(n), 查询: O(1)O(n)
水塘抽样初始化: O(1), 查询: O(n)O(1)

其中 n 为数组长度。水塘抽样方法在空间受限场景下更优，是推荐解法。

相关题目