Random Pick with Weight

Implement a class Solution that samples indices at random according to weights.

• Solution(w) receives a non-negative integer array w; w[i] is the weight of index i.
• pickIndex() returns a random index i with probability w[i] / sum(w).

In particular, an index with weight 0 must never be returned. Each test is a sequence of calls: the first creates the object, and each later one calls pickIndex. Since results are random, the judge accepts any index whose weight is positive, so make sure the distribution is right yourself.

Example 1

Example 2

Example 3

Constraints

• 1 <= w.length <= 104
• 0 <= w[i] <= 105, and at least one weight is positive
• pickIndex is called at most 104 times

Pattern Path · Binary Search · problem 62 of 169

Start simple

Expand each index i into w[i] copies in a list and pick a uniform element (O(sum(w)) memory), or scan the prefix sums linearly on every pick (O(n) per call).

The key idea

Think of the weights as segments laid end to end: index 0 owns the integers [0, w[0]), index 1 owns [w[0], w[0] + w[1]), and so on. The ends of these segments are the prefix sums prefix[i] = w[0] + ... + w[i], and index i owns exactly w[i] of the total integers. So if we draw r uniformly from 0 .. total - 1, the owner of r is picked with probability w[i] / total, which is exactly what we want.

The owner of r is the first index whose prefix sum is strictly greater than r. Because weights are non-negative, the prefix sums never decrease, so this is a binary search (an upper-bound search). A zero-weight index has an empty segment, its prefix sum equals the previous one, and the strict comparison skips over it, so it can never be chosen.

The constructor does O(n) work once, and each pick is O(log n).

Complexity

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

Python

import bisect
import random
from itertools import accumulate

class Solution:
    def __init__(self, w):
        # prefix[i] is the end of index i's segment on the ruler
        self.prefix = list(accumulate(w))
        self.total = self.prefix[-1]

    def pickIndex(self):
        r = random.randrange(self.total)  # uniform in [0, total)
        # first index whose prefix sum is strictly greater than r
        return bisect.bisect_right(self.prefix, r)

JavaScript

class Solution {
  constructor(w) {
    // prefix[i] is the end of index i's segment on the ruler
    this.prefix = [];
    let sum = 0;
    for (const weight of w) {
      sum += weight;
      this.prefix.push(sum);
    }
    this.total = sum;
  }

  pickIndex() {
    const r = Math.floor(Math.random() * this.total); // uniform in [0, total)
    // first index whose prefix sum is strictly greater than r
    let lo = 0, hi = this.prefix.length - 1;
    while (lo < hi) {
      const mid = (lo + hi) >> 1;
      if (this.prefix[mid] > r) hi = mid;
      else lo = mid + 1;
    }
    return lo;
  }
}

What the interviewer is checking

It checks whether you can turn a probability statement into prefix sums plus binary search, and whether your boundary choice (strictly greater vs greater-or-equal) gives each index exactly its share.

Mistakes to avoid

• Drawing r from 1 .. total and then searching for the first prefix greater than r: the first index loses one unit of weight and the last index is out of range.
• Using a lower-bound search with r in [0, total), which can return a zero-weight index or shift every segment by one.
• Recomputing prefix sums on every call instead of once in the constructor.
• Using a floating-point random number and comparing with integer boundaries carelessly.

Follow-ups to expect

• What if weights change between picks? (Fenwick tree over the weights.)
• Make each pick O(1). (Walker's alias method.)
• Pick k distinct indices by weight. (Weighted reservoir sampling.)

In real software

Load balancers routing traffic in proportion to server capacity, A/B testing frameworks splitting users by percentage, and loot tables in games all sample by weight with prefix sums and a binary search.

Practice next