Random Pick with Weight

Problem Statement

Given an array w of positive weights, implement pickIndex() which returns an index with probability proportional to w[index].

Example

Input:  w = [1,3]
Output: index 0 ~25% of the time, index 1 ~75%.

Approach

Build prefix sums so index i owns the range (prefix[i-1], prefix[i]]. Pick a random integer target in [1, total] and binary-search (lower bound) the first prefix sum >= target — that index is chosen with the correct probability.

Complexity

• Time: O(n) build, O(log n) per pick
• Space: O(n)

Solution Code

Python

import bisect, random
class Solution:
    def __init__(self, w):
        self.prefix = []
        total = 0
        for x in w:
            total += x
            self.prefix.append(total)
        self.total = total
    def pickIndex(self):
        target = random.randint(1, self.total)
        return bisect.bisect_left(self.prefix, target)

Java

class Solution {
    int[] prefix; int total;
    Random rand = new Random();
    Solution(int[] w) {
        prefix = new int[w.length];
        int sum = 0;
        for (int i = 0; i < w.length; i++) { sum += w[i]; prefix[i] = sum; }
        total = sum;
    }
    int pickIndex() {
        int target = rand.nextInt(total) + 1;   // 1..total
        int lo = 0, hi = prefix.length - 1;
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (prefix[mid] < target) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
}

C++

class Solution {
    vector<int> prefix; int total;
public:
    Solution(vector<int>& w) {
        int sum = 0;
        for (int x : w) { sum += x; prefix.push_back(sum); }
        total = sum;
    }
    int pickIndex() {
        int target = rand() % total + 1;         // 1..total
        int lo = 0, hi = (int)prefix.size() - 1;
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (prefix[mid] < target) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
};

Key Takeaways

• Weighted sampling reduces to a lower-bound search over prefix sums.
• Pick the target in [1, total] so each weight owns the right-sized band.
• Reference: LeetCode 528 (Random Pick with Weight).