给定一个正整数数组 w
,其中 w[i]
代表下标 i
的权重(下标从 0
开始),请写一个函数 pickIndex
,它可以随机地获取下标 i
,选取下标 i
的概率与 w[i]
成正比。
例如,对于 w = [1, 3]
,挑选下标 0
的概率为 1 / (1 + 3) = 0.25
(即,25%),而选取下标 1
的概率为 3 / (1 + 3) = 0.75
(即,75%)。
也就是说,选取下标 i
的概率为 w[i] / sum(w)
。
示例 1:
输入: inputs = ["Solution","pickIndex"] inputs = [[[1]],[]] 输出: [null,0] 解释: Solution solution = new Solution([1]); solution.pickIndex(); // 返回 0,因为数组中只有一个元素,所以唯一的选择是返回下标 0。
示例 2:
输入: inputs = ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"] inputs = [[[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] ...... 诸若此类。
提示:
1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex
将被调用不超过10000
次
注意:本题与主站 528 题相同: https://leetcode-cn.com/problems/random-pick-with-weight/
“前缀和 + 二分查找”。
class Solution:
def __init__(self, w: List[int]):
n = len(w)
self.presum = [0] * (n + 1)
for i in range(n):
self.presum[i + 1] = self.presum[i] + w[i]
def pickIndex(self) -> int:
n = len(self.presum)
x = random.randint(1, self.presum[-1])
left, right = 0, n - 2
while left < right:
mid = (left + right) >> 1
if self.presum[mid + 1] >= x:
right = mid
else:
left = mid + 1
return left
# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()
class Solution {
private int[] presum;
public Solution(int[] w) {
int n = w.length;
presum = new int[n + 1];
for (int i = 0; i < n; ++i) {
presum[i + 1] = presum[i] + w[i];
}
}
public int pickIndex() {
int n = presum.length;
int x = (int) (Math.random() * presum[n - 1]) + 1;
int left = 0, right = n - 2;
while (left < right) {
int mid = (left + right) >> 1;
if (presum[mid + 1] >= x) {
right = mid;
} else {
left = mid + 1;
}
}
return left;
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(w);
* int param_1 = obj.pickIndex();
*/
class Solution {
public:
vector<int> presum;
Solution(vector<int>& w) {
int n = w.size();
presum.resize(n + 1);
for (int i = 0; i < n; ++i) presum[i + 1] = presum[i] + w[i];
}
int pickIndex() {
int n = presum.size();
int x = rand() % presum[n - 1] + 1;
int left = 0, right = n - 2;
while (left < right)
{
int mid = left + right >> 1;
if (presum[mid + 1] >= x) right = mid;
else left = mid + 1;
}
return left;
}
};
/**
* Your Solution object will be instantiated and called as such:
* Solution* obj = new Solution(w);
* int param_1 = obj->pickIndex();
*/
type Solution struct {
presum []int
}
func Constructor(w []int) Solution {
n := len(w)
pre := make([]int, n+1)
for i := 0; i < n; i++ {
pre[i+1] = pre[i] + w[i]
}
return Solution{pre}
}
func (this *Solution) PickIndex() int {
n := len(this.presum)
x := rand.Intn(this.presum[n-1]) + 1
left, right := 0, n-2
for left < right {
mid := (left + right) >> 1
if this.presum[mid+1] >= x {
right = mid
} else {
left = mid + 1
}
}
return left
}
/**
* Your Solution object will be instantiated and called as such:
* obj := Constructor(w);
* param_1 := obj.PickIndex();
*/