Map Uniform Draws to Indices for Weighted Random Selection
Company: Google
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Onsite
An index of a list of positive integer weights is to be chosen at random, with each index chosen with probability proportional to its weight. In an interview this is usually a class that is built once from the weights and then asked for many picks. Here the randomness is supplied as input so that the result can be checked: each draw is an integer that a caller chose uniformly from `0` to `sum(weights) - 1`, and your function maps every draw to the index it selects.
### Function Signature
```python
def weighted_picks(weights: list[int], draws: list[int]) -> list[int]:
```
### Rules
- Lay the indices out in order along consecutive integer positions starting at `0`: index `0` occupies the first `weights[0]` positions, index `1` the next `weights[1]` positions, and so on, up to position `sum(weights) - 1`.
- A draw `r` selects the index that occupies position `r`. Because every position is equally likely, index `i` is selected with probability `weights[i] / sum(weights)`.
- Return a list holding the selected index for each draw, in the same order as `draws`.
### Constraints
- `1 <= len(weights) <= 10^5`
- `1 <= weights[i] <= 10^4`, so `sum(weights) <= 10^9`, which fits in a 32-bit signed integer
- `1 <= len(draws) <= 10^5`
- `0 <= draws[j] <= sum(weights) - 1`
- Draws may repeat and are in no particular order.
### Examples
**Example 1**
```text
Input: weights = [1, 3, 2], draws = [0, 1, 3, 4, 5]
Output: [0, 1, 1, 2, 2]
```
Index 0 occupies position 0, index 1 occupies positions 1 to 3, and index 2 occupies positions 4 and 5.
**Example 2**
```text
Input: weights = [5], draws = [0, 4, 2]
Output: [0, 0, 0]
```
A single index occupies every position.
**Example 3**
```text
Input: weights = [2, 1, 4, 1], draws = [7, 2, 1, 3, 6]
Output: [3, 1, 0, 2, 2]
```
Index 0 occupies positions 0 and 1, index 1 position 2, index 2 positions 3 to 6, and index 3 position 7.
Overview: A coding question that makes weighted random selection checkable: each index owns a block of positions proportional to its weight, and every supplied uniform draw must be mapped to the index that owns that position. It tests preprocessing the weights for fast lookups over many draws and exact handling of block boundaries.
You are given a list of positive integer `weights`. An index of the list is to be chosen at random, with each index chosen with probability proportional to its weight. In an interview this is usually a class that is built once from the weights and then asked for many picks. Here the randomness is supplied as input so that the result can be checked: each entry of `draws` is an integer that a caller chose uniformly from `0` to `sum(weights) - 1`, and your function maps every draw to the index it selects.
Implement `weighted_picks(weights, draws)` following these rules:
- Lay the indices out in order along consecutive integer positions starting at `0`: index `0` occupies the first `weights[0]` positions, index `1` the next `weights[1]` positions, and so on, up to position `sum(weights) - 1`.
- A draw `r` selects the index that occupies position `r`. Because every position is equally likely, index `i` is selected with probability `weights[i] / sum(weights)`.
- Return a list holding the selected index for each draw, in the same order as `draws`. The result has exactly `len(draws)` entries, each in the range `0` to `len(weights) - 1`.
Every value involved, including `sum(weights)`, is at most `10^9` and fits in a 32-bit signed integer (no value exceeds `2^31 - 1`).
### Constraints
- `1 <= len(weights) <= 10^5`
- `1 <= weights[i] <= 10^4`, so `sum(weights) <= 10^9`, which fits in a 32-bit signed integer
- `1 <= len(draws) <= 10^5`
- `0 <= draws[j] <= sum(weights) - 1`
- Draws may repeat and are in no particular order.
### Example 1
```text
Input: weights = [1, 3, 2], draws = [0, 1, 3, 4, 5]
Output: [0, 1, 1, 2, 2]
```
Index 0 occupies position 0, index 1 occupies positions 1 to 3, and index 2 occupies positions 4 and 5.
### Example 2
```text
Input: weights = [2, 1, 4, 1], draws = [7, 2, 1, 3, 6]
Output: [3, 1, 0, 2, 2]
```
Index 0 occupies positions 0 and 1, index 1 position 2, index 2 positions 3 to 6, and index 3 position 7. The output follows the order of `draws`.
Constraints
- 1 <= len(weights) <= 10^5
- 1 <= weights[i] <= 10^4, so sum(weights) <= 10^9, which fits in a 32-bit signed integer
- 1 <= len(draws) <= 10^5
- 0 <= draws[j] <= sum(weights) - 1
- Draws may repeat and are in no particular order.
Examples
Input: ([1], [0])
Expected Output: [0]
Explanation: Minimum valid input: a single weight of 1 and a single draw at position 0.
Input: ([1, 3, 2], [0, 1, 3, 4, 5])
Expected Output: [0, 1, 1, 2, 2]
Explanation: Source Example 1: index 0 occupies position 0, index 1 positions 1 to 3, index 2 positions 4 and 5.
Hints
- Index i covers one contiguous block of positions. Work out where that block starts and ends using only the weights that come before it.
- There can be 10^5 weights and 10^5 draws, so examining every weight for every draw is too slow. What could you compute once from the weights and reuse for every draw?
- Check the edges: which index owns the position just after one index's block ends, and which owns the last position of that block?