Resolve a Weighted Random Pick

Interview Problem Interview Medium Prefix Sum Array Binary Search

Checking your account…

Loading the interactive Interview workspace.If it does not appear, the problem and learning material remain readable, but browser execution is unavailable.Reload Interview workspace

Problem

Implement Solution.pickIndex(weights, draw). weights contains positive integers and draw is an injected integer from 1 through sum(weights), representing the random sample. Return the first index whose cumulative weight is at least draw. This deterministic contract tests the same prefix-sum lookup used by Random Pick with Weight.

Starter code

class Solution:
    def pickIndex(self, weights, draw):
        pass

Test cases

middle-bucket

{
  "args": [
    [
      1,
      3,
      2
    ],
    4
  ]
}

Expected: 1

Wizard outline

Step 1: Initialize Solution.pickIndex
Replace the empty starter with the first real state owned by Solution.pickIndex. A small, named state is easier to verify than a complete algorithm. Establish it before adding the branch or loop that changes it.
Step 2: Assemble the primary transition
Extend the initialized state with the next contiguous part of the popular solution. The transition explains how one input element or operation changes the state; boundaries are easier to reason about after this invariant is visible.
Step 3: Pass the Next Boundary case
Complete the readable core algorithm for one representative Interview case. Build inclusive prefix weights and binary-search the first cumulative total that reaches a one-based draw.
Step 4: Harden the Middle Bucket boundary
Repair the reviewed boundary and pass the complete submission contract. Index i owns the integer interval after prefix[i-1] through prefix[i]. The first prefix total reaching draw is therefore exactly the unique weighted interval containing the injected sample.

Footguns and prerequisites

A one-based draw uses a greater-than-or-equal boundary; mixing zero-based sampling with this contract causes off-by-one errors.
arrays strings two pointers sliding window

Reviewed references

LeetCode 528 — Random Pick with Weight

Practice prerequisites

Build Prefix Aggregates(opens in a new tab)
Build Prefix Aggregates isolates after consuming i values, the final aggregate equals sum(values[:i]) and the list has i + 1 entries. That focused state discipline is required when implementing resolve weighted random pick as a complete Interview Problem.

Recommended approach and implementation

Accumulate inclusive prefix totals, then lower-bound search for the first total greater than or equal to draw.

Why it works: Index i owns the integer interval after prefix[i-1] through prefix[i]. The first prefix total reaching draw is therefore exactly the unique weighted interval containing the injected sample.

class Solution:
    def pickIndex(self, weights, draw):
        """
        Checkpoint 1: initialize the state owned by this Interview contract.
        Checkpoint 2: assemble the primary transition without hiding the boundary.
        """
        prefix = []
        total = 0
        for weight in weights:
            total += weight
            prefix.append(total)
        left, right = 0, len(prefix) - 1
        while left < right:
            middle = (left + right) // 2
            if prefix[middle] < draw:
                left = middle + 1
            else:
                right = middle
        return left