【Leetcode】528. Random Pick with Weight

题目地址：

https://leetcode.com/problems/random-pick-with-weight/

给定一个长$n$的数组$A$，$A[i]$表示$i$的重复次数，要求设计一个类，可以在$0\sim n-1$这$n$个数里，随机挑选数，使得$x$的出现概率恰好等于$A[i]/{\sum }_{i}A[i]$。

思路是利用前缀和，开一个新数组$p$，其长也为$n$，并且满足$p[0]=A[0],p[i]-p[i-1]=A[i]$。挑选数的过程可以这样做，先利用均匀分布随机变量产生一个$0\sim A[n-1]-1$的随机数$x$，然后找到第一个下标$k$使得$A[k]>x$，则这次挑选出的数即为$k$。显然$0$被选择的概率就是$$\frac{p[0]}{p[n-1]}=\frac{A[0]}{{\sum }_{i}A[i]}$$对于别的数$c$被选择的概率就是$$\frac{p[c]-p[c-1]}{p[n-1]}=\frac{A[c]}{{\sum }_{i}A[i]}$$符合条件。而找到第一个$A[k]>x$的$k$，可以通过二分法得到。代码如下：

class Solution {
 public:
  vector<int> p;

  Solution(vector<int>& w) {
    srand(time(nullptr));
    p.resize(w.size());
    p[0] = w[0];
    for (int i = 1; i < w.size(); i++)
      p[i] = p[i - 1] + w[i];
  }

  int pickIndex() {
    int idx = rand() % p.back();
    int l = 0, r = p.size() - 1;
    while (l < r) {
      int mid = l + (r - l >> 1);
      if (p[mid] > idx) r = mid;
      else l = mid + 1;
    }
    return l;
  }
};

预处理时间复杂度$O(n)$，每次询问时间$O(\log n)$，空间$O(n)$。