Leetcode 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).

URL: https://leetcode.com/problems/remove-k-digits/

My answer:

The problem wants me to simulate picking a random index where each index has a specific weight. So my first thought was to use a random number generator. The Java class Random has the class method nextInt(num) that can generate a random number from 0 to (num - 1). However, every generated number in that range has the same probability of being picked.

This problem was tricky to approach because of this, but the trick to solving the problem is to think about counting numbers. If I had 10 objects and I counted all of them, I would be assigning each object a unique number between 1 and 10 inclusive. But what if I decided to count the objects a different way? I could count the 10 objects as {1,2,3,4,6,7,8,9,10,11}. If I were to generate a random number between 1 and 11, a 5 or 6 being generated would point to the index 4. Therefore, index 4 now has a higher chance of being chosen which simulates a weighted probability. This logic is used to solve the problem.

My solution:

class Solution {
    int[] rPick;
    public Solution(int[] w) {
        rPick = new int[w.length];
        rPick[0] = w[0];
        for (int i = 1; i < w.length; i++) { 
            rPick[i] = w[i] + rPick[i-1];
        }
    }
    
    public int pickIndex() {
        Random rand = new Random();
        int randInt = rand.nextInt(rPick[rPick.length-1]);
        randInt += 1;
        
        //linear search solution
        // for (int i = 0; i < rPick.length; i++) {
        //     if (randInt <= rPick[i]) {
        //         return i;
        //     }
        // }
        
        //binary search
        int p = 0, r = rPick.length - 1;
        while (p < r) {
            int q = p + (r-p)/2;
            if (randInt == rPick[q]) {
                return q;
            }
            else if (randInt > rPick[q]) {
                p = q + 1;
            }
            else {
                r = q;
            }
        }
        return p;
    }
}

/**
 * Your Solution object will be instantiated and called as such:
 * Solution obj = new Solution(w);
 * int param_1 = obj.pickIndex();
 */

Explanation:

Line 4-5 create the array rPick to hold the numbers used to add up the counts(weights) of the indices.

Line 6-8 is the for loop that fills rPick with the sum of counts so far by adding the current count to the previous count. If every index has equal weight, then a 1 would be added to every previous value and the result would be a consecutive sequence of numbers.

Line 12-14 generates a random number from 0 to the maximum value in the rPick array minus 1. A one is added to the result to adjust the number for the problem.

Line 16-21 searches for the index using a linear search. If the randomly generated number is less than the value at the current index, then the current index corresponds to the random number. The index is then returned. This method works, but is the slower option for searching.

Line 24-36 implements the search using a modified binary search algorithm. The index is the target, so the numbers to search are between 0 to rPick.length - 1.

Line 25 is the loops condition for the binary search, that will terminate once p == r. The important thing to note is that this modified algorithm will always result in p being equal to r.

Line 27 returns the index if the random number is equal to the value at the index.

Line 30 checks if the random number is greater than the value at index q. If so, then search the array [q+1,r]. Since q cannot be an answer, it is excluded from future searches.

Line 33 executes if the random number is less than the value at index q. If so, then the array [p,q] is searched. The reason for keeping q in the search, is because the index q can still be the answer if the value there is no value at a index less than q that is also greater than the generated value.

Line 37 executes when the loop finally terminates at p==r. The index p(or index r) is the correct index and is returned.