[LeetCode] 398. Random Pick Index

Given an array of integers 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.

Note:
The array size can be very large. Solution that uses too much extra space will not pass the judge.

Example:

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

// pick(3) should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
solution.pick(3);

// pick(1) should return 0. Since in the array only nums[0] is equal to 1.
solution.pick(1);

Thought process:
The question doesn't have a required time complexity on pick, but rather it doesn't allow using too much extra space. The idea is to store the given array as is. Iterate through the array. Use a variable count to keep track of the number of occurrences of target. And use count as an upper bound for random.

Solution:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

class Solution { private int[] reservoir; private Random random; public Solution(int[] nums) { reservoir = nums; random = new Random(); } public int pick(int target) { int result = -1; int count = 0; for (int i = 0; i < reservoir.length; i++) { if (reservoir[i] != target) { continue; } count++; if (random.nextInt(count) == 0) { result = i; } } return result; } } /** * Your Solution object will be instantiated and called as such: * Solution obj = new Solution(nums); * int param_1 = obj.pick(target); */

Time complexity:
pick(): O(n).