399/670

399. Random Pick with Weight

Medium

You receive an array w of positive integer weights and an array draws of integers. The classic interview task is to build a pickIndex() that returns index i with probability w[i] / sum(w). The standard construction: lay the weights end to end on a number line so index i owns the half-open bucket [prefix[i-1], prefix[i]) of length w[i] (where prefix is the running sum and prefix[-1] = 0), draw a uniform random integer r in [0, sum(w)), and report the bucket that contains r.

So results can be checked, the random draws are handed to you instead of generated: every draws[j] satisfies 0 <= draws[j] < sum(w). Return an array of the same length as draws where entry j is the index whose bucket contains draws[j] — that is, the smallest index i with prefix[i] > draws[j].

Example 1:

Input: w = [1, 3], draws = [0, 1, 2, 3]

Output: [0, 1, 1, 1]

Explanation: The running sums are [1, 4], so index 0 owns [0, 1) and index 1 owns [1, 4). Draw 0 lands in index 0's bucket; draws 1, 2, and 3 all land in index 1's — exactly the 1-in-4 vs 3-in-4 split the weights describe.

Example 2:

Input: w = [2, 5, 3], draws = [4, 9, 0, 6]

Output: [1, 2, 0, 1]

Explanation: Running sums [2, 7, 10] carve [0, 10) into [0, 2) → 0, [2, 7) → 1, [7, 10) → 2. Draw 4 falls in [2, 7), draw 9 in [7, 10), draw 0 in [0, 2), and draw 6 in [2, 7).

Constraints:

  • 1 ≤ w.length ≤ 10⁴
  • 1 ≤ w[i] ≤ 10⁵
  • 1 ≤ draws.length ≤ 10⁴
  • 0 ≤ draws[j] < w[0] + w[1] + ... + w[n-1]

Hints:

Place the weights end to end: index i owns a segment of length w[i] starting where segment i−1 ended. A draw r belongs to index i exactly when the running sum first exceeds r at i. A simple scan that subtracts weights from r until it goes negative answers one draw in O(n).

The running sums form a strictly increasing array (weights are positive). "Smallest index whose prefix sum exceeds r" is a classic binary-search predicate — precompute the prefix array once and answer every draw in O(log n) with an upper-bound search.

▶ Run checks these sample cases. Submit also runs hidden edge cases.

Input: w = [1, 3], draws = [0, 1, 2, 3]

Expected output: [0, 1, 1, 1]