Random Pick Index

The goal here is to implement the Solution class that initializes with an integer array nums and has a method pick, which returns a random index of a given target number from the array.

Here’s the Python implementation:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
import random

class Solution:
    def __init__(self, nums: List[int]):
        self.index_map = collections.defaultdict(list)
        for index, num in enumerate(nums):
            self.index_map[num].append(index)

    def pick(self, target: int) -> int:
        indexes = self.index_map[target]
        return random.choice(indexes)
  • In the constructor, we build a dictionary (index_map) that maps each unique number in the array to the list of its indexes in the array.
  • In the pick method, we retrieve the list of indexes for the given target and randomly pick one using random.choice.

This implementation ensures that if there are multiple valid indexes for the given target, each index will have an equal probability of being returned.