398.随机数索引

题目描述

给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。

注意：
数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。

示例:

int[] nums = new int[] {1,2,3,3,3};
Solution solution = new Solution(nums);

// pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
solution.pick(3);

// pick(1) 应该返回 0。因为只有nums[0]等于1。
solution.pick(1);

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/random-pick-index
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

解题思路

本题需要将nums数组中等于target的值的下标返回，因为nums数组中可能存在多个与target相等的元素，所以当查到多个相等的值时，要能够实现同概率返回符合条件的下标，这题的难点就在于随机返回，且返回的下标的值要是随机的。

思路一：用一个列表来专门存储所有符合题意的下标，再从这个列表中随机的返回一个。通过调用random库中的randint函数实现随机返回下标。这需要开辟额外的辅助空间来存储对应的下标。

思路二：所谓的蓄水池算法，就是每查到一个符合题意的下标，该下标就有与之前记录下标相等的概率被修改，最后返回的一个res的值也因此是所以符合题意下标中的等概率返回的其中一个。

代码

方法1：

class Solution:

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

    def pick(self, target: int) -> int:
        ans=[]
        for i in range(len(self.nums)):
            if self.nums[i]==target:
                ans.append(i)
        index=random.randint(0,len(ans)-1)
        return ans[index]

方法2：

class Solution:

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

    def pick(self, target: int) -> int:
        time,res=0,-1
        for i in range(len(self.nums)):
            if self.nums[i]==target:
                time+=1
                if random.randint(1,time)==1:
                    res=i
        return res

        对于方法二的理解：

        假如输入的是[1,2,3,3,3],要查找的目标值是3

        初始时，time=0,res=-1

        当查到第一个3时，time=1,res=2（第一个3的下标）

        当查到第二个3时，time=2,res如果要更新为3（第二个3的下标），就有1/2的概率

        当查到最后一个3时，time=3,res如果要更新为4(最后一个3的下标)，就有1/3的概率

        最后的3被扫描完后，以后就没有符合题意的3，这时候，res=2的概率就是：从头到尾res都保持不变，没有更新，那么概率就是：1/2  * 2/3 = 1/3,这里面的1/2是在time=2时，没有随机到1，则res保持不变，2/3是指在time=3时，没有随机到1，res保持不变，1,2,3三个数字没有随机到1的概率自然是2/3,所以 res=2概率为1/3；

        同理，当res=3,即第二次被修改为3，要想最后输出的还是3，就要保证最后一次随机时不被修改，那么最后一次不被修改的概率也为2/3,但要想成为3还要满足第一次被修改，第一次被修改的概率为1/2，那么最后输出3的概率也为1/2 * 2/3 = 1/3；

        对于res=4的结果，只要最后一次被修改为4就可以，而不用关心前面几次res存的是哪个下标，所以res=4的最终概率为1/3

        综上，对于每一个符合题意的下标，最终输出的概率是相等的，这也就是这种算法的巧妙之处。