Random Pick with Weight
Asked at Goldman Sachs, LinkedIn
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.
Asked At
| Company | Difficulty | |
|---|---|---|
| Goldman Sachs | Medium | View all Goldman Sachs questions → |
| Medium | View all LinkedIn questions → |
How to Think About It
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.
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.
Precompute prefix sums in the constructor. For each pick, draw r uniformly from 1..total and binary search for the first prefix sum >= r.
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.
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).
What Trips People Up in Real Interviews
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.
Solution Code
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)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 is described in detail above, including step-by-step walkthroughs, complexity analysis, and solution code in Python. Scroll up to the "Optimal Approach" section.
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.