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:

1
2
3
4
5
6
7
8
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);

这种随机数的题目很少做。这个题目用到了蓄水池抽样算法(Reservoir Sampling). wikipedia里有详细说明。

For sample size 1:

Suppose we see a sequence of items, one at a time. We want to keep a single item in memory, and we want it to be selected at random from the sequence. If we know the total number of items (n), then the solution is easy: select an index i between 1 and n with equal probability, and keep the i-th element. The problem is that we do not always know n in advance. A possible solution is the following:

  • Keep the first item in memory.
  • When the i-th item arrives (for i>1):
    • with probability 1/i, keep the new item (discard the old one)
    • with probability 1-1/i, keep the old item (ignore the new one)

So:

  • when there is only one item, it is kept with probability 1;
  • when there are 2 items, each of them is kept with probability 1/2;
  • when there are 3 items, the third item is kept with probability 1/3, and each of the previous 2 items is also kept with probability (1/2)(1-1/3) = (1/2)(2/3) = 1/3;
  • by induction, it is easy to prove that when there are n items, each item is kept with probability 1/n.

For sample k from n:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
public List<Integer> sample(List<Integer> list, int k) {
    final List<Integer> samples = new ArrayList<Integer>(k);
    int count = 0;
    final Random random = new Random();
    for (Integer item : list) {
        if (count < k) {
            samples.add(item);
        } else {
            // http://en.wikipedia.org/wiki/Reservoir_sampling
            // In effect, for all i, the ith element of S is chosen to be included in the reservoir with probability
            // k/i.
            int randomPos = random.nextInt(count);
            if (randomPos < k) {
                samples.set(randomPos, item);
            }
        }
        count++;
    }
    return samples;
}

本题正是要求取一个。代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
public class Solution {
    int[] nums;
    Random rnd;
    public Solution(int[] nums) {
        this.nums = nums;
        this.rnd = new Random();
    }
    public int pick(int target) {
        int result = -1;
        int count = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] != target)
                continue;
            if (rnd.nextInt(++count) == 0)
                result = i;
        }
        return result;
    }
}

Ref: