MediumArrayMathBinary SearchPrefix SumRandomizedUpdated Sep 2026

528. Random Pick with Weight

Asked at Goldman Sachs, LinkedIn

Difficulty

Medium

Time

O(n) to build, O(log n) per pick

Space

O(n)

Pattern

Array

Companies

2

Problem

You are given an array w of positive integers. Each position i has a weight w[i]. You must build a class with a method pickIndex() that returns a random index.

The catch: the pick is not uniform. Index i must come back with probability w[i] / sum(w). So a bigger weight means a bigger chance. For example, with w = [1, 3], index 0 should appear about 25% of the time and index 1 about 75% of the time.

Because the result is random, many different outputs are accepted. What matters is that over many calls the frequencies match the weights.

Weights as slices of a number line

w1301index 0 has weight 1, index 1 has weight 3number line 1..4123401231 lands on index 0; 2, 3, 4 land on index 1 — that is1/4 vs 3/4
For w = [1, 3] the total is 4, so index 0 owns 1 unit and index 1 owns 3 units of the line.

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","pickIndex"] [[[1]],[]]
Output[null,0]
WhyThere is only one index, so every call must return 0. Its probability is 1/1 = 100%.
Example 2
Input["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"] [[[1,3]],[],[],[],[],[]]
Output[null,1,1,1,1,0]
WhyIndex 1 has weight 3 out of a total of 4, so it should show up roughly 3 times out of 4. This run happened to return 1 four times and 0 once. Other mixes like [null,1,1,1,0,0] are also valid.

Asked At

2 companies
CompanyDifficulty
Goldman SachsMediumView all Goldman Sachs questions →
LinkedInMediumView all LinkedIn questions →

How to Think About It

1

Naive approach: expand the array so index i appears w[i] times, then pick uniformly. Weights can be up to 10^5 each, so that array can be enormous.

2

Key insight: lay the weights end to end on a number line. Index i owns the interval (prefix[i-1], prefix[i]]. A uniform random point on [1, total] falls into interval i with exactly the right probability.

3

Precompute prefix sums in the constructor. For each pick, draw r uniformly from 1..total and binary search for the first prefix sum >= r.

4

Walkthrough for w = [1,3]: prefix [1,4]. r = 1 -> index 0; r = 2, 3, 4 -> index 1. So index 1 has probability 3/4.

5

Watch the boundary convention: with r in [1, total] use "first prefix >= r"; with r in [0, total) use "first prefix > r".

Optimal Approach

Constructor: build prefix where prefix[i] = w[0] + ... + w[i]; total = prefix[-1].

pickIndex():
Step 1: r = random integer in [1, total].
Step 2: Binary search for the smallest i with prefix[i] >= r.
Step 3: Return i.

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

Prefix sums plus binary search

prefix1401prefix[0] = 1, prefix[1] = 1 + 3 = 4, and total = 4pick r in 1..412340123r = 1 gives index 0; r = 2, 3 or 4 gives index 1r = randint(1, total)i = first prefix[i] >= rreturn i
Build prefix sums once, then map each random point to the first prefix that reaches it.
ComplexityTime O(n) to build, O(log n) per pickSpace O(n)The cost comes from storing the prefix sums and binary searching them on each pick.

What Trips People Up in Real Interviews

1

Off-by-one on the random range. Drawing from [0, total] inclusive gives index 0 one extra unit of probability.

2

Linear scan of the prefix array per pick. Correct, but O(n) — binary search is expected.

3

Using floating-point division to normalize weights. Integer prefix sums avoid precision issues entirely.

4

Rebuilding the prefix array inside pickIndex. It should be computed once in the constructor.

Solution Code

O(n) to build, O(log n) per pick time
import random
from bisect import bisect_left

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):
        r = random.randint(1, self.total)
        return bisect_left(self.prefix, r)

Pro at DSA?

Test your skills with a real FAANG-style mock interview.

Start a Mock Interview →

Frequently Asked Questions

What is the Random Pick with Weight problem?

Random Pick with Weight gives you an array of positive weights and asks you to implement `pickIndex`, which returns index `i` with probability `w[i] / sum(w)`. It combines prefix sums with binary search — a pattern that shows up in load balancers, sampling, and A/B test bucketing.

How do you solve Random Pick with Weight?

The optimal approach runs in O(n) to build, O(log n) per pick time and O(n) space. Scroll up to "Optimal Approach" for the step-by-step walkthrough, diagrams, and code in Python, C++ and Java.

What companies ask Random Pick with Weight?

Random Pick with Weight is asked at Goldman Sachs, LinkedIn. It is a medium difficulty problem.

What are common mistakes on Random Pick with Weight?
  • Off-by-one on the random range. Drawing from `[0, total]` inclusive gives index 0 one extra unit of probability.
  • Linear scan of the prefix array per pick. Correct, but `O(n)` — binary search is expected.
  • Using floating-point division to normalize weights. Integer prefix sums avoid precision issues entirely.
  • Rebuilding the prefix array inside `pickIndex`. It should be computed once in the constructor.