Skip to main content

528. Random Pick with Weight

medium

Sample an index proportionally to a weight array. The textbook 'prefix sum + binary search' move that powers weighted random sampling everywhere from ads to game loot tables.

By Sam K., Founder, InterviewChamp.AI · Last verified

Problem

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith index. You need to implement the function pickIndex(), which randomly picks an index in the range [0, w.length - 1] (inclusive) and returns it. The probability of picking an index i is w[i] / sum(w).

Constraints

  • 1 <= w.length <= 10^4
  • 1 <= w[i] <= 10^5
  • pickIndex will be called at most 10^4 times.

Examples

Example 1

Input
Solution([1]); pickIndex()
Output
0

Explanation: The only choice is to return 0 since there is only one element in w.

Example 2

Input
Solution([1,3]); pickIndex(); pickIndex(); pickIndex(); pickIndex(); pickIndex()
Output
[0,1,1,1,0]

Explanation: Index 0 is picked with probability 1/4, index 1 with probability 3/4.

Solve it now

Free. No sign-up. Python and JavaScript run instantly in your browser.

Output

Press Run or Cmd+Enter to execute

Hints

Progressive — try the first before opening the next.

Hint 1

Build prefix sums of w so that prefix[i] represents the cumulative weight from 0 through i.

Hint 2

Pick a uniform target in (0, total]. The answer is the smallest index i with prefix[i] >= target.

Hint 3

That's a lower-bound binary search over the prefix array — O(log n) per pick.

Solution approach

Reveal approach

Precompute prefix sums in the constructor: prefix[i] = w[0] + w[1] + ... + w[i]. Each pickIndex() call generates a uniform target in (0, total] where total = prefix[n - 1], then uses lower-bound binary search to find the smallest i with prefix[i] >= target. The intuition: index i 'owns' the half-open interval (prefix[i - 1], prefix[i]], which has length w[i], so it's selected with probability w[i] / total — exactly the requested distribution. Binary search standard template: lo = 0, hi = n - 1; while lo < hi, mid = lo + (hi - lo) / 2; if prefix[mid] < target, lo = mid + 1; else hi = mid. Return lo. Constructor is O(n) time and space; each pick is O(log n). The pattern (cumulative distribution + binary search) is the canonical inverse-transform-sampling implementation and shows up in dozens of domains.

Complexity

Time
O(n) constructor, O(log n) per pick
Space
O(n)

Related patterns

  • binary-search
  • prefix-sum

Related problems

Asked at

Companies reported asking this problem (sourced from public Glassdoor, Blind, and Levels.fyi interview posts).

  • Google
  • Amazon
  • Meta
  • Apple

More Binary Search practice problems

See all Binary Search problems →