LeetCode 398. Random Pick Index

Post by ailswan Aug. 12, 2024

398. Random Pick Index

Problem Statement

link: LeetCode.cn LeetCode Given an integer array nums with possible duplicates, randomly output the index of a given target number. You can assume that the given target number must exist in the array.

Implement the Solution class:

Solution(int[] nums) Initializes the object with the array nums. int pick(int target) Picks a random index i from nums where nums[i] == target. If there are multiple valid i’s, then each index should have an equal probability of returning.

Example:

Input: ["Solution", "pick", "pick", "pick"][[[1, 2, 3, 3, 3]], [3], [1], [3]] Output: [null, 4, 0, 2]

Constraints: 1 <= nums.length <= 2 * 104 -231 <= nums[i] <= 231 - 1 target is an integer from nums. At most 104 calls will be made to pick.

Solution Approach

Store the indices of each number in a dictionary, and then randomly select an index from the list of indices corresponding to the target number.

Algorithm

  1. Store Indices: Use a dictionary to map each number in the array to a list of its indices.
  2. Random Selection: When picking, retrieve the list of indices for the target number.
  3. Return Random Index: Randomly select and return an index from the list using a random choice function.

Implement

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

  def pick(self, target: int) -> int:
      return choice(self.pos[target])