528. Random Pick With Weight

/**

 * 528. Random Pick with Weight

 * You are given an array of positive integers w where w[i] describes the weight of ith index (0-indexed).

 * We need to call the function pickIndex() which randomly returns an integer in the

 * range [0, w.length – 1]. pickIndex() should return the integer proportional to its weight in the w array.

 * For example, for w = [1, 3], the probability of picking the index 0 is 1 / (1 + 3) = 0.25 (i.e 25%)

 * while the probability of picking the index 1 is 3 / (1 + 3) = 0.75 (i.e 75%).

 * More formally, the probability of picking index i is w[i] / sum(w).

 * Example 1:

 * Input

 * [“Solution”,”pickIndex”]

 * [[[1]],[]]

 * Output

 * [null,0]

 * Explanation

 * Solution solution = new Solution([1]);

 * solution.pickIndex(); // return 0. Since there is only one single element on the array the only option is to return the first element.

 */

See solution to the problem on github here: https://github.com/zcoderz/leetcode/blob/main/src/main/java/frequent/medium/RandomPickByWeight.java

This question is working with probability distributions and is different from a typical interview question and hence i have included it here.

The question is asking to pick an index based on its weight as determined by value / total array sum. The algorithm is :

1. Calculate aggregate ways up to an index starting from the first index
2. Use a random number generator to generate a random number between 0 and the total sum of weights. 
3. Run a binary search in the weight array to find the index in the array representing the sum denoted by the random number. 

Here is the code for the solution:

int[] aggWeights;
int sum;
public RandomPickByWeight(int[] w) {
    //add the aggregated weights into a new array.
    //note that this array will be sorted as weights are positive
    aggWeights = new int[w.length];
    int sum = 0;
    for (int i = 0; i < w.length; i++) {
        sum += w[i];
        aggWeights[i] = sum;
    }
    this.sum = sum;
}

public int pickIndex() {
    //get a random number and search for its index in the array
    double rand = Math.random() * sum;
    int low = 0; int hi= aggWeights.length-1;

    while (low < hi) {
        int mid = low + ((hi-low) / 2); //the conversion from double to int will round the double down to int
        if (aggWeights[mid] < rand) {
            //important to move low one ahead of the mid instead of modifying condition to move head down.
            //that's because the above logic rounds down the mid and we could stay in a loop
            //if we dont advance low.
            low = mid+1;
        } else {
            hi = mid;
        }
    }
    return hi;
}