LeetCode-398 随机数索引

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/random-pick-index

题目描述

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

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

示例:

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);

解题思路

很明显这个题的场景是数据太大，内存无法全部加载，所以需要想办法压缩数据。

思路有两个:

一个是压缩数据，因为存在重复数据，所以可以合并数据并且记录重复次数，这样初始化时间复杂度为O(n), pick时间复杂度为O(1).

一个是使用数学规律，并不需要加载数据，所以初始化复杂度为O(1)，但是在pick的时候记录遇到target的次数cnt，然后取一个[0,cnt)的随机数，如果这个随机数为0 则取这个索引，可以证明每个索引的概率是相同的。

证明：假设有k个target，那么在第cnt次取到索引的概率是

代码展示

哈希表暴力法：

class Solution {
public:
    unordered_map<int, vector<int>> m_Nums;
    Solution(vector<int>& nums) {
        for(int i = 0; i < nums.size(); i++)
        {
            m_Nums[nums[i]].push_back(i);
        }
    }

    int pick(int target) {
        int iRand = rand() % m_Nums[target].size();
        return m_Nums[target][iRand];
    }
};

/**
 * Your Solution object will be instantiated and called as such:
 * Solution* obj = new Solution(nums);
 * int param_1 = obj->pick(target);
 */

水池抽样法：

class Solution {
public:
    vector<int> m_Nums;
    Solution(vector<int>& nums)
    :m_Nums(nums)
    {
    }

    int pick(int target) {
        int iCnt = 0, iIndex;
        for(int i = 0; i < m_Nums.size(); i++)
        {
            if(target == m_Nums[i])
            {
                iCnt++;
                if(rand() % iCnt == 0)
                    iIndex = i;
            }
        }
        return iIndex;
    }
};

/**
 * Your Solution object will be instantiated and called as such:
 * Solution* obj = new Solution(nums);
 * int param_1 = obj->pick(target);
 */

运行结果