Return index with probability proportional to its weight

By | February 11, 2015
Share the joy
  •  
  •  
  •  
  •  
  •  
  •  

This question is from careercup. link
Problem description:
Write a function that receives a stream of pairs: id (some unique integer) + weight (positive real number).
The stream ends with the special pair {0,0}.
The function should return one the unique ids at random, with probability proportional to its weight (compared to the total weights sum when stream has ended).
Stream size and weights sum are unknown in advance, and you are allowed O(1) memory consumption.
Example: If this was the stream: {1,2},{2,2},{3,4},{4,10}, then id 2 would be returned with probability 1/9.

Analysis:
Assume the input from [0,W0], [1,W1],…[n,Wn], the selected index is i. Let totalWeight=sum(W0,…,Wn).
Now consider a new input [n+1,Wn+1], the probability to keep index I is totalWeight / (totalWeight + Wn+1), and probbability to have a new index n+1 is Wn+1 / (totalWeight + Wn+1).

Instead of a input stream, I used int[][] array for short.

public static int getRandom(int[][] input) {
    double totalWeight = input[0][1];
    int result = 0;
    for (int i = 1; i < input.length; i++) {
        if (input[i][0] == 0 && input[i][1] == 0) {
            break;
        }
        double newTotalWeight = totalWeight + input[i][1];
        double newChance = ((double) input[i][1]) / newTotalWeight;
        if (Math.random() <= newChance) {
            result = i;
        }
        totalWeight = newTotalWeight;
    }
    return result;
}

public static void main(String[] args) {
    int[][] input = { { 1, 4 }, { 2, 2 }, { 3, 2 }, { 0, 0 } };
    int result = getRandom(input);
    System.out.println(result);
    /*
     * Following runs for 1000 times to check the answer correctness.
     *
     * int[] test = new int[3]; for (int i = 0; i < 1000; i++) {
     * test[getRandom(input)]++; } for (int i = 0; i < test.length; i++) {
     * System.out.print(test[i] + “\t”); }
     */
}