leetcode-528.按权重随机选择
数学问题-随机与取样
题目详情
给你一个 下标从 0 开始 的正整数数组 w
,其中 w[i]
代表第 i
个下标的权重。
请你实现一个函数 pickIndex
,它可以 随机地 从范围 [0, w.length - 1]
内(含 0 和 w.length - 1)选出并返回一个下标。选取下标 i 的 概率 为 w[i] / sum(w)
。
例如,对于 w = [1, 3],挑选下标 0 的概率为 1 / (1 + 3) = 0.25 (即,25%),而选取下标 1 的概率为 3 / (1 + 3) = 0.75(即,75%)。
示例1:
输入:
["Solution","pickIndex"]
[[[1]],[]]
输出:
[null,0]
解释:
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0,因为数组中只有一个元素,所以唯一的选择是返回下标 0。
示例2:
输入:
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
输出:
[null,1,1,1,1,0]
解释:
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1,返回下标 1,返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0,返回下标 0,返回该下标概率为 1/4 。
由于这是一个随机问题,允许多个答案,因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。
思路:
先使用 partial_sum
求前缀和(即到每个位置为止之前所有数字的和)
每当需要采样时,我们可以先随机产生一个数字,然后使用二分法
查找其在前缀和中的位置,以模拟加权采样的过程。
这里借鉴力扣题解中的一个例子解释:
假设我们有数组w: [1, 2, 3, 4], 那么这个数组的的和为 1 + 2 + 3 + 4 = 10. 对应的我们得到 index {0,1,2,3} 的概率为 {1/10, 2/10, 3/10, 4/10}。现在我们要返回 {0,1,2,3} 中的随意一个index,但是我们要保证pickIndex()函数所得到这个index的概率是根据以上的权重来的。
我们把每个数(即权重)拆为1等分
我的代码:
class Solution
{
vector<int> sums;
public:
Solution(vector<int>& w)
{
//sums存储对应下标的前缀和
//如sums[0] = w[0] sums[1] = sums[0] + w[1] ....
partial_sum(w.begin(), w.end(), back_inserter(sums));
}
/*等价于
sums.push_back(W[0]);
for(int i = 1; i < w.size(); ++i)
sums.push_back(sums.back() + w[i]);
*/
int pickIndex()
{
int pos = (rand() % sums.back()) + 1; //产生0~sums.back()之间的随机数
//二分法查找pos在sums区间里的位置,减去begin的差值即为pos的下标即为所求
return lower_bound(sums.begin(), sums.end(), pos) - sums.begin();
}
};
/**
* Your Solution object will be instantiated and called as such:
* Solution* obj = new Solution(w);
* int param_1 = obj->pickIndex();
*/