algoadvance

Leetcode 528. Random Pick with Weight

Problem Statement

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

• Write a function pickIndex which randomly picks an index in proportion to its weight.

Here is a specific example:

• w = [1, 3] means that the chance of picking index 0 is 1 / (1+3) = 0.25 (weight of index 0 divided by total weight), and the chance of picking index 1 is 3 / (1+3) = 0.75.

Clarifying Questions

1. Is the w array size fixed?
   • For this problem, we’ll assume that the weights array is immutable during the lifetime of our class instance.
2. Do the weights always sum to a manageable number, ensuring probabilities are non-zero?
   • Yes, all weights are positive integers, hence their combined sum will be positive, guaranteeing non-zero probabilities.
3. How should the pickIndex method be tested?
   • The effectiveness of the pickIndex method can be tested by repeatedly calling it many times (e.g., 10,000 times) and checking if the frequencies of the indices match the expected probabilities.

Strategy

1. Prefix Sum Array: First, we compute a prefix sum array such that each element at index i contains the cumulative sum of weights up to index i. This helps us to determine intervals corresponding to each index.
2. Random Number Generation: Generate a random number in the range [1, total_weight] where total_weight is the sum of all weights.
3. Binary Search: Use binary search to find the appropriate index corresponding to the generated random number.

Code

import java.util.Random;

class Solution {
    private int[] prefixSums;
    private int totalWeight;
    private Random random;

    public Solution(int[] w) {
        prefixSums = new int[w.length];
        totalWeight = 0;
        random = new Random();
        
        for (int i = 0; i < w.length; i++) {
            totalWeight += w[i];
            prefixSums[i] = totalWeight;
        }
    }

    public int pickIndex() {
        int target = random.nextInt(totalWeight) + 1;
        // Binary search to find the right index
        int low = 0, high = prefixSums.length - 1;
        
        while (low < high) {
            int mid = low + (high - low) / 2;
            if (prefixSums[mid] < target) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }
        
        return low;
    }

    public static void main(String[] args) {
        int[] w = {1, 3}; // Example initialization
        Solution solution = new Solution(w);
        
        // Testing the distribution
        int n = 10000;
        int[] count = new int[w.length];
        for (int i = 0; i < n; i++) {
            count[solution.pickIndex()]++;
        }
        
        // Display the result
        for (int i = 0; i < w.length; i++) {
            System.out.println("Index " + i + " picked " + (count[i] * 100.0 / n) + "% times.");
        }
    }
}

Time Complexity

• Initialization (Solution constructor): O(n), where n is the number of elements in the weights array w. This is because we iterate through the weights array to compute the prefix sums.
• Pick Index (pickIndex method): O(log(n)), where n is the number of elements in the weights array w. This is due to the binary search algorithm used to determine the index.

This solution is efficient for both initialization and picking an index, making it suitable for use cases that require high-frequency random picks.

Cut your prep time in half and DOMINATE your interview with AlgoAdvance AI

• Got blindsided by a question you didn’t expect?

• Spend too much time studying?

• Or simply don’t have the time to go over all 3000 questions?