题目难度: 中等
今天继续更新 Leetcode 的剑指 Offer(专项突击版)系列, 大家在公众号 算法精选 里回复
剑指offer2
就能看到该系列当前连载的所有文章了, 记得关注哦~
题目描述
给定一个正整数数组 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 次
题目思考
- 如何设计随机算法?
- 如何优化时间复杂度?
解决方案
思路
- 分析题目, 要考虑每个下标的权重, 我们显然不能直接随机取
[0,n)
的下标, 这样会导致每个下标取到的概率相同, 如何解决呢? - 根据题目例子, 我们可以发现, 每个下标的取值概率等于其权重除以权重总和
- 利用这一点, 假设权重总和是 total, 我们随机取
[0,total)
之间的某个数字, 然后判断它落在了哪个下标的范围内, 这样就把下标权重考虑在内了 - 举个例子, 假设数组 w 是
[3,2,4]
, 那么总权重和就是3+2+4=9
- 下标 0 的权重是 3, 它负责的范围是[0,2]
- 下标 1 的权重是 2, 它负责的范围是[3,4]
- 下标 2 的权重是 4, 它负责的范围是[5,8]
- 我们随机取
[0,9)
之间的某个数字, 它落在哪个下标范围就返回哪个下标, 这样就保证了每个下标的取值概率等于其权重除以权重总和 - 有了思路后, 如何写出代码呢?
- 不难发现每个下标的负责范围的起点就是其前缀和, 即
[0,3,5,9]
, 注意最后的 9 代表了总和, 不属于某个下标起点 - 所以我们可以在初始化时求出原始数组对应的前缀和数组 sms, 最后一个元素
sms[-1]
就是总和 - 然后 pickIndex 时随机取
[0,sms[-1])
之间的数字 x, 并从头开始遍历前缀和数组, 找到第一个满足sms[i]<=x<sms[i+1]
的下标 i, 即为考虑权重后随机选取的下标 - 不过上述查找过程需要线性遍历, 时间复杂度是 O(N), 能否继续优化呢?
- 注意到题目给出的是正整数数组, 所以前缀和是单调递增的, 我们可以利用二分查找来优化
- 这里我们能够完美利用 Python 提供的右二分 bisect_right, 它的语义是:
- 如果数字存在于查找数组, 返回对应下标加 1
- 如果数字不存在, 则返回首个大于该数字的下标或数组长度(如果已有数字都不大于查找数字时)
- 这样右二分查找到的下标正是所需下标的下一个下标, 同样拿上面例子举例:
- 假如查找数字是 0, 则右二分返回下标 1, 所需下标是 0
- 假如查找数字是 1, 则右二分返回下标 1, 所需下标是 0
- 假如查找数字是 2, 则右二分返回下标 1, 所需下标是 0
- 假如查找数字是 3, 则右二分返回下标 2, 所需下标是 1
- 假如查找数字是 4, 则右二分返回下标 2, 所需下标是 1
- …
- 大家如果感兴趣的话也可以自己尝试实现这个二分查找的过程, 基本类似前面的题目Leetcode 剑指 Offer II 068.搜索插入位置, 只需要稍作改动, target 存在时返回对应下标加 1, 而不是下标本身
- 下面代码中有详细的注释, 方便大家理解
复杂度
- 时间复杂度 O(logN): 每次 pickIndex 使用二分查找, 时间复杂度是 O(logN)
- 空间复杂度 O(N): 需要额外存储前缀和数组
代码
class Solution:
def __init__(self, w: List[int]):
# 前缀和+二分查找
self.sms = [0]
for x in w:
# 统计前缀和数组
self.sms.append(x + self.sms[-1])
def pickIndex(self) -> int:
# 随机选择[0,数字总和)之间的一个数字x
x = random.randrange(0, self.sms[-1])
# 二分查找当前数字落在了哪个下标的范围内
# 这里使用右二分bisect_right, 且二分得到的下标要减1
return bisect.bisect_right(self.sms, x) - 1
大家可以在下面这些地方找到我~😊
我的公众号: 算法精选, 欢迎大家扫码关注~😊