按权重随机选择

今日题目

今日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%）。

---

题目解读

读第一遍的时候，我的表情就像上面这样。
严重怀疑自己的理解能力。
还好后来看懂了，浅浅解读一下

输入数据： w:List[int]
(即，元素为正整数的列表w)

已知条件：
1.w的索引下标从0开始；
2.w中每个元素的值代表该元素对应索引下标的权重

题目要求：
设计一个 pickIndex 函数，按权重占比为概率返回一个w的索引下标

这其实是一个广泛见于抽奖、道具抽取的应用需求
由于每个奖品价值不同，那么商家肯定不会给你等概率抽到啦。贵的奖品权重低，抽到的概率极低

举个例子
有一个奖品列表 idx=[0,1,2,3]
每个奖品被抽中的权重，组成的权重列表w=[4,1,2,3]
那么，显而易见
idx[0]被抽中的概率是4/(1+2+5+2)=4/10=0.4

---

算法思想

此题的重点在于如何以权重占比为概率去随机选出一个数
我们计算一下权重总和t，构造一个范围1~t的列表r

由上图易知，从列表r中随机一个数，若随机数为1~4，那么对应的就是选出了idx[0]；若随机数为5，那么对应的就是选出了idx[1] …

现在问题转化为如何表现这种对应关系
我们构造一个前缀和列表psum

那么这个时候就很容易啦，从r中随机出数字n，在psum中查找有序插入n的索引i（若有相同元素则插入在左边）

例如，我们随机到了6，那么在psum中其应该插入的索引i为2

关于查找算法，由于psum本身就是有序的，因此采用折半查找法效率更高

代码实现

class Solution:

    def __init__(self, w: List[int]):
        # 生成前缀和列表
        self.psum = list(accumulate(w))

    def pickIndex(self) -> int:
        # 生成1~t的随机数，t为权重总和
        n = random.randint(1, self.psum[-1])
        # 前缀和列表中查找n的索引
        return bisect_left(self.psum, n)

如果要在我们自己的编辑器里实现，则要完善一下引包以及main函数的实现

from typing import List
from itertools import accumulate
import random
from bisect import bisect_left

class Solution:

    def __init__(self, w: List[int]):
        # 生成前缀和列表
        self.psum = list(accumulate(w))

    def pickIndex(self) -> int:
        # 生成1~t的随机数，t为权重总和
        n = random.randint(1, self.psum[-1])
        # 前缀和列表中查找n的索引
        return bisect_left(self.psum, n)
        
def main():
    w = [4, 1, 3, 2]
    print(Solution(w).pickIndex())

if __name__=='__main__':
     main()

---

另外还看到了一个很牛逼的O(1)复杂度的算法。过几天续写在这里❀❀❀