leetcode398 随机数索引-蓄水池算法

leetcode398 随机数索引

题目要求

给你一个可能含有 重复元素 的整数数组 nums ，请你随机输出给定的目标数字 target 的索引。你可以假设给定的数字一定存在于数组中。

实现 Solution 类：

Solution(int[] nums) 用数组 nums 初始化对象。
int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i 。如果存在多个有效的索引，则每个索引的返回概率应当相等。

解题思路

该问题的核心有两个：

1. 问题的规模可能非常大，以至于无法完全存储下来
2. 有多个下标可以返回，但是要求等概率。

算法

蓄水池算法。
参考链接：https://www.cnblogs.com/snowInPluto/p/5996269.html

解决该问题的核心思路就是，遍历到第k个满足要求的样本，就要按照 \(1/k\) 的概率去选择它。

代码

class Solution:

    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        c = 0
        ans = 0
        for i, v in enumerate(self.nums):
            if v == target:
                c += 1
                if randrange(c) == 0:
                    ans = i
        return ans

这里的randrange方法是random包中的方法，\(random.randrange ([start,] stop [,step])\)

这里我们就赋值了一个参数，也就是stop。该函数会在\([0, stop-1]\)范围内等概率选择一个样本，如果等于0的话，概率就是\(1/c\)，这是满足要求的。