AlgoMaster Logo

Random Pick with Weight

medium7 min readUpdated June 23, 2026

Understanding the Problem

We need to build a system that picks random indices, but not uniformly. Each index has a weight, and heavier indices should be picked more often. If index 0 has weight 1 and index 1 has weight 3, then index 1 should be picked 3 times as often as index 0.

One way to model this is a number line. If the weights are [1, 3, 2], the total is 6. Index 0 owns 1 unit of that line, index 1 owns 3 units, and index 2 owns 2 units. A uniform random point on the line lands in each index's segment with probability proportional to that index's weight.

The work is in mapping a random number to the correct index efficiently. Scanning the weights on every call is one option, but prefix sums combined with binary search do it faster.

Key Constraints:

  • 1 <= w.length <= 10^4 → Up to 10,000 weights. O(n) preprocessing in the constructor is fine.
  • 1 <= w[i] <= 10^5 → All weights are positive, so prefix sums are strictly increasing and every index has a non-zero probability. The maximum total weight is 10^4 * 10^5 = 10^9, which still fits in a 32-bit signed integer.
  • pickIndex will be called at most 10^4 times → With O(n) per call, that is up to 10^8 comparisons total. O(log n) per call brings it down to roughly 10^4 * 14 comparisons.

Approach 1: Linear Scan with Prefix Sum

Intuition

Build a prefix sum array from the weights, generate a random number between 1 and the total weight, then scan the prefix sums to find which index that random number falls into.

The prefix sum array turns the problem into a range lookup. If the weights are [1, 3, 2], the prefix sums are [1, 4, 6]. Index 0 covers the range [1, 1], index 1 covers [2, 4], and index 2 covers [5, 6]. A random number of 3 falls into index 1's range, so we return 1.

Algorithm

  1. In the constructor, build a prefix sum array from the weights.
  2. In pickIndex(), generate a random integer in the range [1, totalWeight].
  3. Scan through the prefix sum array. Return the first index where prefixSum[i] >= randomTarget.

Example Walkthrough

1Prefix sums built from weights [1, 3, 2]. Random target = 3
0
1
i
1
4
2
6
1/4

Code

The linear scan ignores a property of the prefix sum array: it is always in increasing order. Binary search exploits that to find the target range in O(log n) instead of O(n).

Approach 2: Prefix Sum + Binary Search

Intuition

Because all weights are positive, the prefix sum array is strictly increasing. That ordering lets binary search replace the linear scan.

Rather than checking each prefix sum one by one, binary search finds the smallest index where prefixSum[i] >= target. This is the lower bound pattern: locate the first element not less than a target value in a sorted array.

Algorithm

  1. In the constructor, build a prefix sum array from the weights. Store the total weight.
  2. In pickIndex():
    • Generate a random integer target in the range [1, totalWeight].
    • Use binary search to find the smallest index i where prefixSum[i] >= target.
    • Return that index.

Example Walkthrough

1Prefix sums: [1, 4, 6]. Binary search for first >= 5. left=0, right=3
0
1
left
1
4
2
6
search range
1/5

Code

Binary search makes each pick O(log n) with O(n) preprocessing. If picks vastly outnumber the array size, the alias method drives query time down to O(1) after a one-time O(n) build.

Approach 3: Alias Method (O(1) per pick)

Intuition

The alias method (Vose's variant) restructures the distribution so that a single pick needs only one random index and one coin flip, both O(1).

The idea is to flatten an uneven distribution into n equal-height columns. Each column belongs to one primary index but can also store an "alias" index that fills the rest of the column. Scale every probability by n so the average column height is exactly 1. Indices lighter than average donate their leftover space to indices heavier than average, and each donation is recorded as an alias.

After the tables are built, a pick chooses a column uniformly (that is the O(1) part), then flips a biased coin to decide between the column's primary index and its alias.

Algorithm

  1. In the constructor, compute scaled[i] = w[i] * n / total, so the values sum to n and average 1.
  2. Split indices into a small list (scaled < 1) and a large list (scaled >= 1).
  3. While both lists are non-empty:
    • Take an index s from small and l from large.
    • Set prob[s] = scaled[s] and alias[s] = l (column s is scaled[s] filled by s, the rest by l).
    • Reduce l's remaining space: scaled[l] = scaled[l] - (1 - scaled[s]).
    • Move l to small if scaled[l] < 1, otherwise keep it in large.
  4. Any index left in either list gets prob = 1 (it fills its own column completely).
  5. In pickIndex(), choose a uniform column col in [0, n). Draw a uniform r in [0, 1). Return col if r < prob[col], otherwise return alias[col].

Example Walkthrough

1Weights [1, 3, 2], total=6, n=3. scaled[i] = w[i]*3/6 = [0.5, 1.5, 1.0]
0
0.5
1
1.5
2
1
1/6

The columns now reproduce the original distribution. Index 1 is reached from its own column (prob 1.0) plus the alias halves of columns 0 and 2, giving total mass 1 + 0.5 + 0.5 = 2 out of n = 3 columns, which is 2/3 = 3/6, matching w[1]/total.

Code

For this problem's constraints (n up to 10^4, at most 10^4 picks), binary search is the common choice because it is shorter to write and the O(log n) factor is negligible. The alias method matters when the number of picks dominates, since it removes the per-pick log factor entirely.