给定一个正整数数组 w
,其中 w[i]
代表位置 i
的权重,请写一个函数 pickIndex
,它可以随机地获取位置 i
,选取位置 i
的概率与 w[i]
成正比。
说明:
1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex
将被调用不超过10000
次
示例1:
输入: ["Solution","pickIndex"] [[[1]],[]] 输出: [null,0]
示例2:
输入: ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"] [[[1,3]],[],[],[],[],[]] 输出: [null,0,1,1,1,0]
输入语法说明:
输入是两个列表:调用成员函数名和调用的参数。Solution
的构造函数有一个参数,即数组 w
。pickIndex
没有参数。输入参数是一个列表,即使参数为空,也会输入一个 [] 空列表。
思路:这道题要实现不等概率分布的采样,具体参考这道题:Random Point in Non-overlapping Rectangles 非重叠矩形中的随机点,中心思想是通过map的由小到大的自排序性,把累加和作为key,下标作为value,这样在[0,sum)内随机取值的时候就可以保证是不等概率取值。
参考代码如下:
class Solution {
public:
Solution(vector<int> w) {
this->w = w;
for (auto it : w) {
sum += it;
m.insert({ sum,(int)m.size() });
}
}
int pickIndex() {
int value = rand() % sum;
int idx = m.upper_bound(value)->second;
return idx;
}
private:
vector<int> w;
map<int, int> m;
int sum = 0;
};