Implement weighted random index picker
Company: Meta
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Onsite
Design a class WeightedPicker that is initialized with an array of positive integer weights w and supports pick() that returns an index i with probability proportional to w[i]. Optimize for up to 10^7 picks after preprocessing 10^5 weights. Discuss preprocessing choices (prefix sums vs. alias method), random number generation, numerical stability, and time/space complexity. Extension: support addWeight(i, delta) and pick() with sublinear-time updates and queries.
Quick Answer: Implement weighted random index picker evaluates algorithm design, data structures, correctness, complexity, edge cases, and implementation details in a realistic interview setting. A strong answer states assumptions, handles edge cases, explains trade-offs, and shows how to validate the result clearly.
You are designing the core of a `WeightedPicker` (LC528 "Random Pick with Weight"). Given an array `weights` of positive integers, index `i` should be chosen with probability proportional to `weights[i]`. The standard approach precomputes prefix sums and, to pick, draws a uniform integer `target` in `[1, total]` (where `total = sum(weights)`) and returns the first index whose cumulative prefix sum is `>= target`.
Because randomness cannot be unit-tested directly, this challenge tests the **deterministic core**: instead of generating the random draws internally, the draws are supplied to you. Implement `weightedPick(weights, targets)` that, for each `target` in `targets` (each a 1-based integer in `[1, total]`), returns the index `i` such that `prefix[i-1] < target <= prefix[i]`, found via binary search over the prefix sums. Return the list of picked indices in order.
The interval of length `weights[i]` assigned to index `i` guarantees the pick probability is proportional to `weights[i]`.
Example: `weights = [5, 2, 3]` has prefix sums `[5, 7, 10]`. A draw of `1` lands in `[1,5]` -> index 0; a draw of `6` lands in `[6,7]` -> index 1; a draw of `8` lands in `[8,10]` -> index 2.
Discussion points the interviewer expects: prefix-sum + binary search gives O(n) preprocessing and O(log n) per pick; the alias method gives O(1) per pick after O(n) preprocessing; a Fenwick/BIT supports the `addWeight(i, delta)` extension in O(log n) per update and per pick.
Constraints
- 1 <= weights.length <= 10^5
- 1 <= weights[i] (all weights are positive integers)
- Each target is a 1-based integer in [1, sum(weights)]
- 0 <= targets.length <= 10^7 picks
- sum(weights) fits in a 64-bit integer
Examples
Input: ([1], [1])
Expected Output: [0]
Explanation: One weight owns the whole range [1,1]; every draw returns index 0.
Input: ([1, 3], [1, 2, 3, 4])
Expected Output: [0, 1, 1, 1]
Explanation: Prefix sums [1,4]. Draw 1 -> index 0; draws 2,3,4 fall in (1,4] -> index 1. Index 1 is picked 3/4 of the time, matching weight 3 vs 1.
Hints
- Precompute cumulative prefix sums: prefix[i] = weights[0] + ... + weights[i]. Index i 'owns' the half-open range (prefix[i-1], prefix[i]] of width weights[i].
- For a draw 'target', you want the first index whose prefix sum is >= target. That is a classic lower-bound binary search.
- Be careful with the binary-search invariant: use lo<hi with hi=mid (not mid-1) when prefix[mid] >= target, so you converge on the first qualifying index rather than overshooting.
- Extension: to support addWeight(i, delta) with sublinear updates, replace the static prefix array with a Fenwick tree (BIT) — both the point update and the prefix-sum binary search become O(log n).