528. 按权重随机选择

290 阅读1分钟

题目:
给你一个 下标从 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

}