Problem statement

https://leetcode.com/problems/random-pick-index/

Solution

For each element create list of indexes which has this elements and then just use random sample to choose one of the elements.

Complexity

Time is O(n) for init and O(1) for pick. Space is O(n).

Code

class Solution:
    def __init__(self, nums):
        self.d = defaultdict(list)
        for i, num in enumerate(nums):
            self.d[num] += [i]
        
    def pick(self, target):
        return random.sample(self.d[target], 1)[0]

Remark

There is also reservoir sampling technique but it has O(n) complexity and not good for this problem.