Learnitweb

LeetCode Problem 528: Random Pick with Weight

Problem Overview

You are given an array of positive integers w, where each element represents a weight associated with an index.
Your task is to design a function pickIndex() that randomly returns an index in proportion to its weight.

In other words, an index with a larger weight should be more likely to be chosen.

Example:

Input: 
["Solution", "pickIndex"]
[[[1, 3]], []]

Output: 

[null, 1]

Explanation:

  • The array w = [1, 3]
  • The probability of picking index 0 is 1 / (1 + 3) = 0.25
  • The probability of picking index 1 is 3 / (1 + 3) = 0.75

So, index 1 should be returned about 75% of the time.


Intuition / Approach

This problem can be efficiently solved using the concept of prefix sums combined with binary search.

Step 1: Build a Prefix Sum Array

We first compute a prefix sum array that represents the cumulative weight distribution.
For example, if w = [1, 3, 2], then the prefix sum becomes [1, 4, 6].

This means:

  • Index 0 covers range (0, 1]
  • Index 1 covers range (1, 4]
  • Index 2 covers range (4, 6]

Step 2: Random Number Generation

Generate a random number r in the range [1, totalSum], where totalSum is the sum of all weights.
The goal is to find which index range this random number falls into.

Step 3: Binary Search

We use binary search on the prefix sum array to efficiently find the smallest prefix that is greater than or equal to the random number r.
The index of that prefix is our randomly selected index.


Algorithm Explanation

  1. Preprocessing:
    • Compute prefix sums for the input array w.
    • Store the total sum of weights.
  2. pickIndex():
    • Generate a random number r in the range [1, totalSum].
    • Use binary search to find the first index in the prefix sum array where the prefix is greater than or equal to r.
    • Return that index.

Java Code Implementation

import java.util.Random;

public class RandomPickWithWeight {
    private int[] prefixSums;
    private int totalSum;
    private Random random;

    public RandomPickWithWeight(int[] w) {
        prefixSums = new int[w.length];
        prefixSums[0] = w[0];

        // Build prefix sum array
        for (int i = 1; i < w.length; i++) {
            prefixSums[i] = prefixSums[i - 1] + w[i];
        }

        totalSum = prefixSums[w.length - 1];
        random = new Random();
    }

    public int pickIndex() {
        int target = random.nextInt(totalSum) + 1;  // Random number in [1, totalSum]
        return binarySearch(target);
    }

    // Binary search for first prefix >= target
    private int binarySearch(int target) {
        int left = 0;
        int right = prefixSums.length - 1;

        while (left < right) {
            int mid = left + (right - left) / 2;
            if (prefixSums[mid] < target) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }

        return left;
    }

    // Example usage
    public static void main(String[] args) {
        int[] w = {1, 3, 2};
        RandomPickWithWeight solution = new RandomPickWithWeight(w);

        // Simulate multiple picks
        int[] count = new int[w.length];
        for (int i = 0; i < 10000; i++) {
            int index = solution.pickIndex();
            count[index]++;
        }

        System.out.println("Frequency of picks:");
        for (int i = 0; i < w.length; i++) {
            System.out.println("Index " + i + ": " + count[i]);
        }
    }
}

Complexity Analysis

Time Complexity:

  • Preprocessing: O(n) for building the prefix sum array.
  • pickIndex(): O(log n) for binary search on prefix sums.

Space Complexity:

  • O(n) for storing the prefix sum array.

Alternative Approaches

  1. Linear Search (Simpler but Slower):
    Instead of using binary search, you could iterate linearly through the prefix sum array until you find where the random number fits.
    • Time complexity: O(n) per pickIndex()
    • This approach is simpler but inefficient for large input sizes.
  2. Alias Method (Advanced):
    For very large datasets with frequent queries, the Alias Method can be used to generate weighted random picks in O(1) time after O(n) preprocessing.
    However, it’s more complex to implement and beyond the scope of most interview solutions.