Random Pick with Weight

Picking a random item is easy. But what if some items should come up more often than others? That is the real question here. The interviewer wants to see if you can turn weights into fair chances. And then make each pick fast.

🎯 The Problem

You pick a random index, but heavier indices must come up more often. The rules:

• You get an array of weights. Each weight tells you how heavy one index is.
• Pick an index at random.
• The chance of picking an index must match its weight.
• A heavier index gets picked more.
• The chance of an index is called its probability. That is just how likely it is to happen.

Say the weights are [1, 3]. Index 0 has weight 1. Index 1 has weight 3. The total weight is 4. So index 0 should come up 1 time out of 4. Index 1 should come up 3 times out of 4.

Input:  weights = [1, 3], pick a number r in [0, 4)
Output: an index 0 or 1

Explanation:
  total weight = 1 + 3 = 4
  r in [0, 1)  -> index 0   (chance 1/4)
  r in [1, 4)  -> index 1   (chance 3/4)

To keep this example reproducible, we fix the random draw to r = 2. With r = 2, the answer is index 1.

Here is the idea drawn out. Each weight takes up a slice of a number line. The bigger the weight, the bigger the slice.

🐢 Approach 1: Expanded List (Brute Force)

Build a big list with one copy of each index per unit of weight, then pick a random spot.

The idea:

• For each index, add it to a list as many times as its weight.
• Weight [1, 3] becomes the list [0, 1, 1, 1].
• Index 0 appears once. Index 1 appears three times.

How it works:

• Pick one random spot in this big list.
• The value at that spot is your answer.
• Index 1 fills three of the four spots, so it gets picked three times more often. That is the fairness we wanted.

Why it is weak:

• If one weight is a million, you store a million copies.
• The memory explodes with big weights.
• Fine for tiny weights only.

Here is the expanded-list code for that idea:

import random


class Solution:
    def __init__(self, w):
        self.values = []
        for index, weight in enumerate(w):
            self.values.extend([index] * weight)

    def pickIndex(self):
        return random.choice(self.values)

⚡ Approach 2: Prefix Sum Plus Binary Search (Best)

Store only the slice boundaries, then binary search which slice the draw lands in.

The idea:

• We do not need the giant list. We only need the boundary of each slice.
• Build a running total called a prefix sum. Each spot holds the sum of all weights up to and including that index.
• For weights [1, 3], the prefix sums are [1, 4]. The slice for index 0 ends at 1. The slice for index 1 ends at 4.

How it works:

• Pick a random number r from 0 up to the total weight.
• Find the first prefix sum that is strictly greater than r. That index is the answer.
• The prefix sums are always sorted, because we keep adding positive weights.
• A sorted list means binary search applies. Binary search cuts the search range in half each step.

Why it is fast:

• Setup builds a small array once, in O(n).
• Each pick is one binary search, O(log n).
• Memory stays O(n) no matter how big the weights are.

Here is a dry run with weights [1, 3] and r = 2. We search the prefix sums [1, 4] for the first value above 2.

Steps to Solve

1. Build a prefix sum array from the weights. Each spot holds the total weight so far.
2. The last prefix sum is the total weight.
3. Pick a random number r from 0 up to the total weight, not including the total.
4. Binary search the prefix sums for the first value that is strictly greater than r.
5. Return that index as the chosen pick.

#include <stdio.h>

int prefix[100];  // running totals of the weights
int n;            // how many weights

void build(int weights[], int count) {
    n = count;
    int running = 0;
    for (int i = 0; i < n; i++) {
        running += weights[i];   // add this weight to the total
        prefix[i] = running;     // store the running total
    }
}

int pickIndex(int r) {
    int lo = 0, hi = n - 1;
    while (lo < hi) {                 // shrink the range each loop
        int mid = lo + (hi - lo) / 2;
        if (prefix[mid] <= r) {       // slice ends too early, go right
            lo = mid + 1;
        } else {                      // slice covers r, keep it
            hi = mid;
        }
    }
    return lo;  // first prefix strictly greater than r
}

int main(void) {
    int weights[] = {1, 3};
    build(weights, 2);
    int r = 2;  // fixed draw so the output is reproducible
    printf("%d\n", pickIndex(r));
    return 0;
}

1

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

class Solution {
    vector<int> prefix;  // running totals
public:
    Solution(vector<int>& weights) {
        int running = 0;
        for (int w : weights) {
            running += w;            // add this weight
            prefix.push_back(running);
        }
    }
    int pickIndex(int r) {
        // upper_bound gives the first value strictly greater than r
        auto it = upper_bound(prefix.begin(), prefix.end(), r);
        return (int)(it - prefix.begin());
    }
};

int main() {
    vector<int> weights = {1, 3};
    Solution s(weights);
    int r = 2;  // fixed draw so the output is reproducible
    cout << s.pickIndex(r) << endl;
    return 0;
}

1

public class RandomPick {
    static int[] prefix;  // running totals of the weights

    static void build(int[] weights) {
        prefix = new int[weights.length];
        int running = 0;
        for (int i = 0; i < weights.length; i++) {
            running += weights[i];   // add this weight
            prefix[i] = running;     // store the running total
        }
    }

    static int pickIndex(int r) {
        int lo = 0, hi = prefix.length - 1;
        while (lo < hi) {                  // shrink the range each loop
            int mid = lo + (hi - lo) / 2;
            if (prefix[mid] <= r) {        // slice ends too early, go right
                lo = mid + 1;
            } else {                       // slice covers r, keep it
                hi = mid;
            }
        }
        return lo;  // first prefix strictly greater than r
    }

    public static void main(String[] args) {
        int[] weights = { 1, 3 };
        build(weights);
        int r = 2;  // fixed draw so the output is reproducible
        System.out.println(pickIndex(r));
    }
}

1

import bisect


class Solution:
    def __init__(self, weights):
        self.prefix = []
        running = 0
        for w in weights:
            running += w              # add this weight
            self.prefix.append(running)  # store the running total

    def pick_index(self, r):
        # bisect_right finds the first prefix strictly greater than r
        return bisect.bisect_right(self.prefix, r)


weights = [1, 3]
s = Solution(weights)
r = 2  # fixed draw so the output is reproducible
print(s.pick_index(r))

1

class Solution {
  constructor(weights) {
    this.prefix = [];
    let running = 0;
    for (const w of weights) {
      running += w; // add this weight
      this.prefix.push(running); // store the running total
    }
  }

  pickIndex(r) {
    let lo = 0;
    let hi = this.prefix.length - 1;
    while (lo < hi) {
      // shrink the range each loop
      const mid = lo + Math.floor((hi - lo) / 2);
      if (this.prefix[mid] <= r) {
        lo = mid + 1; // slice ends too early, go right
      } else {
        hi = mid; // slice covers r, keep it
      }
    }
    return lo; // first prefix strictly greater than r
  }
}

const weights = [1, 3];
const s = new Solution(weights);
const r = 2; // fixed draw so the output is reproducible
console.log(s.pickIndex(r));

1

Let us walk through the Python version line by line. The setup loop reads running += w for each weight. This keeps a running total. So after weights [1, 3] the prefix list becomes [1, 4]. We store the running total, not the raw weight, because the boundaries are what we search later.

The pick happens in one line. bisect.bisect_right(self.prefix, r) returns the position where r would go to keep the list sorted, just to the right of any equal values. That position is exactly the first prefix strictly greater than r. With r = 2 and prefix [1, 4], the value 1 is not greater than 2, but 4 is. So the answer is position 1. We use bisect_right and not bisect_left because a draw equal to a boundary must fall into the next slice, matching how we defined the ranges.

⏱️ Time and Space Complexity

The brute force builds a giant list, so it can use a huge amount of memory and slow setup. The prefix sum builds a small array once, then every pick is a binary search. So the setup is O(n) and each pick is O(log n). You trade the giant list for a fast search over a tiny array. That keeps memory at O(n) no matter how big the weights are.

🧩 Key Takeaways

• ✅ A weight becomes a slice on a number line. Bigger weight means bigger slice.
• ✅ Prefix sums give you the boundary of each slice with no giant list.
• ✅ Pick a random number, then binary search for the first prefix greater than it.
• ✅ Use a right-leaning binary search so a draw on a boundary lands in the next slice.
• ✅ Setup is O(n) and each pick is O(log n), with only O(n) memory.

🚀 What’s Next?

• Find the Celebrity
• Two Sum