LCR 071. 按权重随机选择

zhumeili / 2024-03-17 / 原文

题目：给定一个正整数数组 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) 。

题解
这道题考察的主要知识点是前缀和和二分查找
主要分成三个步骤：

求出权重比率
在0到1之间生成随机数
求该随机数对应的最小index

class Solution:
    def __init__(self, w: List[int]):
        self.w = w
        self.index = [i for i in range(len(self.w))]
        total_sum = sum(self.w)
        self.rate = [weight / total_sum for weight in self.w]

    def pickIndex(self) -> int:
        start = 0.0
        random_num = random.random()
        for idx, score in enumerate(self.rate):
            start += score
            if random_num <= start:
                return self.index[idx]

但是这样查找的效率很低，最好是使用二分查找，在python中已经封装好了二分查找，可以直接使用

## 这是目前最优的算法
class Solution:

    def __init__(self, w: List[int]):
        self.w = w 
        self.rank = [w[0]]
        for i in range(1,len(self.w)):
            self.rank.append(self.rank[-1] + self.w[i]) ## 只用了一个循环
        

    def pickIndex(self) -> int:
        r = random.randint(1,self.rank[-1])
        return bisect.bisect_left(self.rank,r)
# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()