题目:
给你一个 下标从 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%)。
来源:力扣(LeetCode) 链接:leetcode.cn/problems/ra… 著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。
算法:
数组前缀和:数组的每个元素是之前所有元素之和。得到如下数组:
----4----------8--------------20
求一个20之内的随机数rand,落在哪个范围则取对应的index,如rand<=4则取index=0
方法一、遍历查找
import "math/rand"
type Solution struct {
weightArr []int
}
func Constructor(w []int) Solution {
s := Solution{
weightArr: make([]int, len(w)),
}
s.weightArr[0] = w[0]
for i := 1; i < len(s.weightArr); i ++ {
s.weightArr[i] = s.weightArr[i - 1] + w[i]
}
return s
}
func (this *Solution) PickIndex() int {
// n= 0~max-1
ans := 0
n := rand.Intn(this.weightArr[len(this.weightArr) - 1]) + 1
for i := 0; i < len(this.weightArr); i ++ {
ans = i
if n <= this.weightArr[i] {
break
}
}
return ans
}
**方法二、二分查找 **
func (this *Solution) PickIndex() int {
// n= 0~max-1
ans := 0
randN := rand.Intn(this.weightArr[len(this.weightArr) - 1]) + 1
l, r := 0, len(this.weightArr) - 1
mid := 0
for l <= r {
mid = (l + r) / 2
if randN < this.weightArr[mid] {
r = mid - 1
ans = mid
} else if this.weightArr[mid] < randN {
l = mid + 1
} else {
ans = mid
break
}
}
return ans
}