Leetcode-398-随机数索引

题目

398. 随机数索引

给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。

注意：
数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。

示例:

int[] nums = new int[] {1,2,3,3,3};
Solution solution = new Solution(nums);

// pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
solution.pick(3);

// pick(1) 应该返回 0。因为只有nums[0]等于1。
solution.pick(1);

解题

方法一(暴力)

遍历一遍数组nums，每个nums的元素都作为一个key，将下标记录到一个哈希表中的对应key所属的集合中，pick时候查找哈希表，在查询出的集合中随机挑选一个返回。

class Solution {
public:
    unordered_map<int, vector<int>> mv;
    Solution(vector<int>& nums) {
        for(int i=0;i<nums.size();i++) {
            if(mv.find(nums[i]) == mv.end()) mv[nums[i]] = {};
            mv[nums[i]].push_back(i);
        }
    }
    
    int pick(int target) {
        vector<int>& e = mv[target];
        return e[rand() % e.size()];
    }
};

• 时间复杂度初始化为$O(n)$，pick为$O(1)$

• 空间复杂度$O(n)$

方法二(蓄水池抽样)

方法一适用于已知样本数量且样本数量不超过电脑内存大小的情况。假如我们事先无法知道样本数量，或者样本数量比较庞大无法完全存放在内存中，那么应当使用方法二。

具体前置学习与证明参考三叶姐姐的文章

【蓄水池抽样】多语言入门「蓄水池抽样」

具体做法为：从前往后处理每个样本，使每个样本成为答案的概率为$\frac{1}{i}$ ，其中 $i$ 为样本编号（编号从 1 开始），最终可以确保每个样本成为答案的概率均为 $\frac{1}{n}$（其中 $n$ 为样本总数）。

证明参考上述链接

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

注意这里的$rand() \% idx$的可能取值为$0,1,...,idx-1$，故$P\{rand()\%idx == 0\}=\frac{1}{idx}$

• 时间复杂度：初始化$O(1)$，pick操作$O(n)$

• 空间复杂度：$O(1)$

相似题目

这道题还有类似题目

382. 链表随机节点