398. 随机数索引

思路：蓄水池抽样
时间复杂度O(n)

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

    def pick(self, target: int) -> int:
        # 对多个相同的元素以相同的概率输出，单个元素直接输出
        import random
        count = 0
        ans = 0
        for i,val in enumerate(self.nums):
            if val==target:
                count+=1
                if random.randrange(0,count)==0: ans = i # 记录索引
        return ans