Design weighted random index picker
Company: Meta
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Take-home Project
Design and implement a data structure that, given an array of positive integer weights w[0..n−1], supports pick() that returns index i with probability proportional to w[i]. Optimize for many pick() calls after a one-time initialization. Specify the API, prove correctness, and analyze time and space. Follow-ups: support point updates to weights; handle large totals without precision loss; compare prefix-sum binary search with the alias method; design tests to detect sampling bias.
Quick Answer: This question evaluates algorithm design and data structure selection for weighted random sampling, along with probability reasoning, correctness proof, time and space complexity analysis, and numerical stability considerations.
Implement the deterministic core of a weighted random index picker (LeetCode 528, Random Pick with Weight).
You are given `weights`, an array of non-negative integer weights, and `rolls`, a list of integer "random rolls". In a real picker, `pick()` draws a uniform integer `t` in `[1, total]` (where `total = sum(weights)`) and returns the index `i` of the bucket that `t` lands in. Here the rolls are supplied explicitly so the logic is testable.
Return an array `result` where `result[k]` is the index chosen for `rolls[k]`. The chosen index for a roll `t` is the smallest index `i` such that the prefix sum `w[0] + w[1] + ... + w[i] >= t`. This places `t` in the half-open interval owned by bucket `i`, so index `i` is selected with probability proportional to `w[i]`.
If a roll is out of range (`t < 1` or `t > total`), or if `weights` is empty, append `-1` for that roll.
The intended approach is a one-time O(n) prefix-sum build, then O(log n) binary search per roll — this is the part interviewers grade. (Follow-ups in the original prompt: point updates via a Fenwick/BIT, the alias method for O(1) sampling, avoiding precision loss on huge totals by sampling integers rather than floats, and chi-square tests to detect sampling bias.)
Constraints
- 1 <= len(weights) <= 10^4 in the original picker; this variant also accepts an empty weights array, for which every roll yields -1.
- 0 <= weights[i]; the original problem guarantees positive weights, but a 0 weight is handled (its bucket is never selected).
- Each roll t is an integer; a roll is valid only when 1 <= t <= sum(weights).
- sum(weights) fits in a 64-bit integer.
Examples
Input: ([1], [1])
Expected Output: [0]
Explanation: Single bucket. prefix = [1], total = 1. Roll 1 -> smallest i with prefix[i] >= 1 is index 0.
Input: ([1, 3], [1, 2, 3, 4])
Expected Output: [0, 1, 1, 1]
Explanation: prefix = [1, 4], total = 4. Roll 1 lands in bucket 0; rolls 2,3,4 (>1, <=4) land in bucket 1. Bucket 1 owns 3/4 of the range, matching its weight 3.
Hints
- Compute the prefix sums of the weights once. prefix[i] is the cumulative weight through index i, and the total is prefix[-1].
- A roll t belongs to bucket i if prefix[i-1] < t <= prefix[i]. That is the smallest i with prefix[i] >= t — a classic lower-bound binary search (bisect_left in Python, lower_bound in C++).
- Guard the edges: if weights is empty (total 0) or the roll is < 1 or > total, return -1 for that roll instead of an index.