Code Implementation

import random
import bisect

class Solution:
    def __init__(self, w):
        self.prefix_sums = []
        prefix = 0
        for weight in w:
            prefix += weight
            self.prefix_sums.append(prefix)
        self.total = prefix

    def pickIndex(self):
        target = random.randint(1, self.total)
        # Find the index where target would fit in prefix_sums
        return bisect.bisect_left(self.prefix_sums, target)
      

#include <vector>
#include <cstdlib>
#include <numeric>
#include <algorithm>

using namespace std;

class Solution {
    vector<int> prefixSums;
    int total;
public:
    Solution(vector<int>& w) {
        prefixSums.resize(w.size());
        partial_sum(w.begin(), w.end(), prefixSums.begin());
        total = prefixSums.back();
    }
    
    int pickIndex() {
        int target = rand() % total + 1;
        return lower_bound(prefixSums.begin(), prefixSums.end(), target) - prefixSums.begin();
    }
};
      

import java.util.*;

class Solution {
    private int[] prefixSums;
    private int total;

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

    public int pickIndex() {
        int target = (int)(Math.random() * total) + 1;
        int left = 0, 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;
    }
}
      

class Solution {
    constructor(w) {
        this.prefixSums = [];
        let prefix = 0;
        for (let i = 0; i < w.length; i++) {
            prefix += w[i];
            this.prefixSums.push(prefix);
        }
        this.total = prefix;
    }

    pickIndex() {
        let target = Math.floor(Math.random() * this.total) + 1;
        // Binary search for the first prefixSum >= target
        let left = 0, right = this.prefixSums.length - 1;
        while (left < right) {
            let mid = Math.floor((left + right) / 2);
            if (this.prefixSums[mid] < target) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        return left;
    }
}
      

Problem Description

The "Random Pick with Weight" problem asks you to implement a data structure that, given a list of positive integers w, allows you to randomly pick an index such that the probability of picking index i is proportional to w[i]. In other words, the higher the weight at an index, the more likely it is to be chosen.

You must implement two methods:

• __init__(self, w): Initializes the object with the given list of weights w.
• pickIndex(self): Randomly returns an index according to the weight distribution.

Key constraints:

• All weights are positive integers.
• Each call to pickIndex must be independent and follow the probability distribution defined by the weights.
• There may be multiple calls to pickIndex after initialization.

Thought Process

At first glance, it might seem natural to simulate the process by creating a large array in which each index appears as many times as its weight, and then randomly picking from this array. For example, if w = [1, 3, 2], you could make an array [0, 1, 1, 1, 2, 2] and pick a random element.

However, this approach is inefficient, especially when weights are large, as it uses a lot of extra space and is slow for large inputs. The key is to find a way to map a random number to an index based on the weights, without explicitly building the expanded array.

The optimized idea is to use prefix sums and binary search. By treating the weights as segments on a number line, we can pick a random number in the total weight range and find which segment it falls into. This approach is both space and time efficient.

Solution Approach

The efficient solution uses prefix sums and binary search to map random numbers to weighted indices. Here’s a step-by-step breakdown:

1. Prefix Sum Construction:
   • Compute the prefix sum array of the weights. Each element at index i in this array represents the sum of all weights up to and including w[i].
   • For example, if w = [2, 5, 3], the prefix sums are [2, 7, 10].
2. Random Number Generation:
   • Generate a random integer between 1 and the total sum of weights (inclusive). This simulates picking a random point on the number line formed by the weights.
3. Binary Search:
   • Use binary search on the prefix sum array to find the first index where the prefix sum is greater than or equal to the random number.
   • This index corresponds to the segment (weight) that contains the random point, and thus is the index to return.

This method is efficient because:

• Prefix sum array is built once during initialization (O(n) time and space).
• Each pickIndex call takes O(log n) time due to binary search.

Example Walkthrough

Let's walk through an example with w = [1, 3, 2]:

1. Prefix sums: [1, 4, 6]
   • Index 0: covers 1 (from 1 to 1)
   • Index 1: covers 3 (from 2 to 4)
   • Index 2: covers 2 (from 5 to 6)
2. Random pick: Suppose we randomly pick the number 5 (between 1 and 6).
   • We search for the smallest prefix sum ≥ 5. The prefix sums are [1, 4, 6].
   • 1 < 5, 4 < 5, 6 ≥ 5 → index 2.
3. Result: We return index 2.
4. Distribution: Over many calls, index 0 will be picked 1/6 of the time, index 1 will be picked 3/6, and index 2 will be picked 2/6, matching their weights.

Time and Space Complexity

Brute-force approach:

• Time: O(total_weight) per pick (because you’d build a large array of size equal to the sum of all weights and pick randomly from it).
• Space: O(total_weight) for the expanded array.

• Initialization: O(n) time and space to build the prefix sum array.
• Each pickIndex: O(log n) time for binary search, O(1) space.

This is a significant improvement, especially for large weights or many calls.

Summary

The "Random Pick with Weight" problem is elegantly solved using prefix sums and binary search. By mapping each weight to a segment on a number line and picking a random point, we can efficiently select indices according to their weights. The approach is both time and space efficient, making it suitable for repeated random picks and large input sizes.

Key insights include:

• Representing weights as cumulative segments (prefix sums)
• Using binary search to quickly map a random number to its corresponding index
• Achieving optimal performance for both initialization and picking