【Leetcode】398. Random Pick Index

题目地址：

https://leetcode.com/problems/random-pick-index/

给定一个数组$A$，要求设计一个数据结构，给定一个数$t$，可以随机等概率返回其在$A$中的下标（中的一个）。题目保证$A$里含$t$。

用一个哈希表存一下每个数出现的下标列表即可。代码如下：

class Solution {
 public:
  unordered_map<int, vector<int>> mp;
  Solution(vector<int>& a) {
    srand(time(nullptr));
    for (int i = 0; i < a.size(); i++) mp[a[i]].push_back(i);
  }

  int pick(int x) {
    auto& v = mp[x];
    return v[rand() % v.size()];
  }
};

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

初始化时间复杂度$O(l_A)$，pick时间复杂度$O(1)$，空间$O(l_A)$。

也可以采用所谓的“蓄水池抽样算法”。参考https://blog.csdn.net/qq_46105170/article/details/108807072。这题里只需考虑值为target的下标即可。代码如下：

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

  int pick(int x) {
    int cnt = 0, idx = 0;
    for (int i = 0; i < a.size(); i++)
      // 只看等于x的数
      if (a[i] == x) {
        cnt++;
        // 以1 / cnt的概率上位
        if (rand() % cnt == 0) idx = i;
      }
    return idx;
  }
};
/**
 * Your Solution object will be instantiated and called as such:
 * Solution* obj = new Solution(nums);
 * int param_1 = obj->pick(target);
 */

时间复杂度$O(n)$，空间$O(1)$。