Generated by Codex with GPT-5

Quick facts

What the problem is really asking

The class receives an array of positive weights, where w[i] describes how likely index i should be picked. Calling pickIndex() should return one index at random, but not uniformly. Index i should be returned with probability:

w[i] / sum(w)

For example, if the weights are [1, 3], index 0 should be picked one quarter of the time and index 1 should be picked three quarters of the time.

The key is to stop thinking of the weights as separate probabilities. Instead, imagine laying them out as consecutive ranges on one number line. With [1, 3], the total weight is 4, so the ranges are:

  • index 0: target values 1
  • index 1: target values 2, 3, 4

If a random target is chosen uniformly from 1 through 4, each target value is equally likely. Index 1 owns three target values, so it naturally gets three times the probability of index 0.

The optimal idea

Build prefix sums once in the constructor.

For weights [2, 5, 3], the prefix sums are:

[2, 7, 10]

These prefix sums define the ending point of each index’s range:

  • index 0 owns targets 1..2
  • index 1 owns targets 3..7
  • index 2 owns targets 8..10

Then each call to pickIndex() does two steps:

  1. Generate a random integer target from 1 to total_weight.
  2. Find the first prefix sum that is greater than or equal to target.

That first prefix sum identifies the range containing the target, which gives the weighted random index.

The lookup should use binary search, not a linear scan. The constructor pays O(n) once to build prefix sums, and each pick costs O(log n).

How to derive this in an interview

Start with a brute-force mental model. If w = [2, 5, 3], one way to make the probabilities work would be to expand the array into:

[0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

Then a uniform random pick from that expanded array returns each index with exactly the right probability.

That proves the idea, but it is too wasteful because the weights can be large. The improvement is to keep only the boundaries of those repeated blocks. Prefix sums store those boundaries compactly:

[2, 7, 10]

Now a random integer target plays the role of a random position in the expanded array, and binary search finds which compressed block contains it.

Python solution

from bisect import bisect_left
from random import Random
from typing import List, Optional


class WeightedIndexPicker:
    """
    Pick array indices with probability proportional to their positive weights.

    The prefix sums compactly represent the same distribution as an expanded
    array that repeats index i exactly weights[i] times, but without paying
    memory proportional to the total weight.
    """

    def __init__(
        self,
        weights: List[int],
        random_generator: Optional[Random] = None,
    ) -> None:
        if not weights:
            raise ValueError("weights must contain at least one value")

        self._prefix_sums: List[int] = []
        running_total = 0

        for weight in weights:
            if weight <= 0:
                raise ValueError("all weights must be positive")

            running_total += weight
            self._prefix_sums.append(running_total)

        self._total_weight = running_total
        self._random = random_generator if random_generator is not None else Random()

    def pick_index(self) -> int:
        """
        Return an index i with probability weights[i] / sum(weights).
        """
        target = self._random.randint(1, self._total_weight)
        return bisect_left(self._prefix_sums, target)


class Solution:
    def __init__(self, w: List[int]):
        self._picker = WeightedIndexPicker(w)

    def pickIndex(self) -> int:
        return self._picker.pick_index()

Why this works

Each index owns exactly w[i] integer target values on the range 1..sum(w). Since every target value is chosen with equal probability, the chance of landing in index i’s range is exactly the size of that range divided by the total number of targets:

w[i] / sum(w)

Prefix sums are enough because they store where each range ends. The first prefix sum greater than or equal to the target is the end of the range containing that target. Since prefix sums are strictly increasing for positive weights, binary search finds that range correctly.

The constructor takes O(n) time and O(n) space. Each call to pickIndex() takes O(log n) time for the binary search and O(1) extra space.

Interview follow-ups

What if pickIndex() is called many more times than the constructor?

This prefix-sum approach is designed for exactly that case. The constructor does the one-time O(n) preprocessing, then each pick is only O(log n). If the interviewer wants even faster repeated picks, discuss the alias method, which can reduce each pick to O(1) after O(n) preprocessing. The tradeoff is that alias tables are harder to implement correctly and easier to explain poorly, so prefix sums are usually the better interview solution unless constant-time sampling is explicitly required.

What if the weights can be updated between picks?

Plain prefix sums are not ideal when updates are frequent, because changing one weight affects every later prefix sum and costs O(n). A Fenwick tree is the natural upgrade. It supports point updates in O(log n), maintains prefix sums dynamically, and can binary-search the tree to find the first prefix sum reaching a random target in O(log n). The probability argument stays the same; only the data structure used to store searchable cumulative weights changes.

What if the weights are floating-point probabilities?

If the values are already probabilities, one option is to build floating-point prefix sums and generate a random float in [0, total). Then search for the first prefix sum greater than the target. This works, but it introduces precision concerns around boundary values. If the original weights can be scaled to integers without losing meaning, integer targets are cleaner. If not, the floating-point version is acceptable as long as the boundary convention is consistent.

What if one weight is extremely large?

The prefix-sum method still works because it never expands the weight into repeated entries. It stores one cumulative integer per index. In Python, integers grow as needed, so overflow is not a practical issue. In fixed-width languages, the total weight should be stored in a 64-bit integer when constraints allow large sums. The time complexity depends on the number of weights, not on the magnitude of the weights.

Why use bisect_left with targets from 1 to total_weight?

Using a one-based target makes the range boundaries easy to reason about. If the prefix sums are [2, 7, 10], target 2 should still map to index 0, and bisect_left returns the first boundary that is at least 2. Target 3 maps to index 1, because the first boundary at least 3 is 7. A zero-based target version also works, but it needs a slightly different comparison convention, usually searching for the first prefix sum strictly greater than the target.