Random Pick with Weight

MediumPrefix Sum · Binary Search · Randomized🧩 Binary Search LeetCode #528 ↗Updated Oct 30, 2026

Problem

You are given positive weights w, one per index. Implement a class Solution(int[] w) whose pickIndex() returns a random index i with probability w[i] / sum(w). Many picks are made after a single construction, so pickIndex should be fast.

Inputw = [1]; pickIndex()Output0 (the only index)

Inputw = [1,3]; many pickIndex()Output0 about 25% of the time, 1 about 75%

1 ≤ w.length ≤ 10⁴ · 1 ≤ w[i] ≤ 10⁵ · ≤ 10⁴ calls to pickIndex

Try it yourself

Java

class Solution {
    public Solution(int[] w) {
        
    }

    public int pickIndex() {
        
    }
}

Your code is saved in this browser.Submit on LeetCode ↗

Key insight

Lay the weights end to end on a number line of length total: index i owns a segment of length w[i]. Draw a uniform ticket in [0, total); the segment it lands in is the answer, and the prefix sums are the sorted segment ends, so binary search finds it.

PatternBinary Search

Instinct Sorted data, a rotated array, or find the min/max X that works where a yes/no test flips only once - halve the search space every step.

Memory hack Find the fence - every candidate reads NO NO NO YES YES, and binary search just walks to where the fence stands.

Open the Binary Search guide →

Approaches

Brute ForceExpanded Ticket Pool

Time O(total) build, O(1) pickSpace O(total)

Literally hand out tickets: put index i into a pool w[i] times and pick a uniformly random slot. Each index is chosen in proportion to how many tickets it holds.

Java

class Solution {
    Random rand = new Random();
    private final int[] pool; // index i appears w[i] times
    public Solution(int[] w) {
        int total = 0;
        for (int x : w) total += x;
        pool = new int[total];
        int k = 0;
        for (int i = 0; i < w.length; i++) {
            for (int c = 0; c < w[i]; c++) pool[k++] = i;
        }
    }
    public int pickIndex() {
        return pool[rand.nextInt(pool.length)];
    }
}

Instant picks, but total can reach 10⁹ — the pool will not fit in memory.

BetterPrefix Sums + Linear Scan

Time O(n) build, O(n) pickSpace O(n)

Instead of storing every ticket, store where each index’s run of tickets ends: prefix[i] = w[0] + … + w[i]. A ticket r in [0, total) belongs to the first index with r < prefix[i], found by scanning.

Java

class Solution {
    Random rand = new Random();
    private final int[] prefix; // prefix[i] = w[0] + ... + w[i]
    public Solution(int[] w) {
        prefix = new int[w.length];
        int sum = 0;
        for (int i = 0; i < w.length; i++) {
            sum += w[i];
            prefix[i] = sum;
        }
    }
    public int pickIndex() {
        int r = rand.nextInt(prefix[prefix.length - 1]); // ticket in [0, total)
        for (int i = 0; i < prefix.length; i++) {
            if (r < prefix[i]) return i; // ticket falls in index i's segment
        }
        return prefix.length - 1; // unreachable
    }
}

Memory drops to O(n), but each pick may scan all 10⁴ entries.

OptimalPrefix Sums + Binary Search

Time O(n) build, O(log n) pickSpace O(n)

The prefix array is strictly increasing (weights are positive), so “first index with prefix[i] > r” is an upper-bound search. Each index i owns exactly w[i] ticket values, so the distribution is exact.

Java

class Solution {
    Random rand = new Random();
    private final int[] prefix; // prefix[i] = w[0] + ... + w[i]
    public Solution(int[] w) {
        prefix = new int[w.length];
        int sum = 0;
        for (int i = 0; i < w.length; i++) {
            sum += w[i];
            prefix[i] = sum;
        }
    }
    public int pickIndex() {
        int r = rand.nextInt(prefix[prefix.length - 1]); // ticket in [0, total)
        int lo = 0, hi = prefix.length - 1;
        while (lo < hi) { // first i with prefix[i] > r
            int mid = lo + (hi - lo) / 2;
            if (prefix[mid] > r) hi = mid;
            else lo = mid + 1;
        }
        return lo;
    }
}

Dry runw = [1]; pickIndex()

lo lowm midhi high changed

w=[1,3,2]
1

4

6

prefix sums, total=6
r=3, lo=0 hi=2
1
lo
4
m
6
hi
4 > 3 → hi=1
lo=0 hi=1
1
lo m
4
hi
6

1 ≤ 3 → lo=1
lo=hi=1
1

4
lo hi
6

pick index 1 ✔

Same tiny memory as the scan with logarithmic picks — the standard answer.

Pro tip

Tickets 1..3 map to index 1 here: exactly w[1] = 3 values, which is why the probabilities come out right.

At a glance

Level	Approach	Time	Space
Brute Force	Expanded Ticket Pool	`O(total) build, O(1) pick`	`O(total)`
Better	Prefix Sums + Linear Scan	`O(n) build, O(n) pick`	`O(n)`
Optimal	Prefix Sums + Binary Search	`O(n) build, O(log n) pick`	`O(n)`

Common pitfalls

Mixing ranges: ticket in [0, total) needs prefix[i] > r, not ≥.
Drawing a double and comparing loosely — off-by-one bias at boundaries.
Recomputing prefix sums inside pickIndex instead of once in the constructor.

Interview tips

Draw the number line with segments of length w[i] before writing code.
Check the boundary: the largest ticket total - 1 must map to the last index.
Mention the alias method for O(1) picks if asked about heavy query loads.

Pattern

Prefix sums + binary search → sampling or locating a position in a sequence of variable-length segments.

My notes saved in this browser