Random Pick with Weight

We will be discussing problem : https://leetcode.com/problems/random-pick-with-weight

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the i^th index.

You need to implement the function pickIndex(), which randomly picks an index in the range [0, w.length – 1] (inclusive) and returns it. The probability of picking an index i is:


w[i] / sum(w)

For example, if w = [1, 3], the probability of picking:

Index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%)
Index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%)

Wait, That Sounds Odd…

Generally, when we implement a random function, it is supposed to be uniform:


if we pick a random number in range 10,
prob(1) = 1/10 = 10%
prob(2) = 1/10 = 10%
prob(3) = 1/10 = 10%
...

But now our question says: picking an index has probability proportional to w[i].

Let’s Take an Example: w = [1,2,3,4]

Where the total sum = 10:

index 0: 1 / 10 = 10%
index 1: 2 / 10 = 20%
index 2: 3 / 10 = 30%
index 3: 4 / 10 = 40%

Interesting — pickIndex() Only Returns Index

So how can we map this weighted behavior into a fair implementation?

Idea: Make It Uniform by Expanding Repeats

If we repeat indices proportional to their weights, we can simulate uniform selection:


Original w = [1, 4]
Expanded pool = [0, 1, 1, 1, 1] // indices of 1 and 4

So index 0 is picked 1 out of 5 times → 20%
Index 1 is picked 4 out of 5 → 80%

Generalized Pool Example


Original w = [1, 2, 3, 4]
Expanded pool = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3]

Now we randomly pick:


random index:   actual value:
0               0
1               1
2               1
3               2
4               2
5               2
6               3
7               3
8               3
9               3

Frequencies:

0 – 10%
1 – 20%
2 – 30%
3 – 40%

Works Well but Memory Heavy

The method above works beautifully for small weights. But if you have:

w = [1, 1000000, 2]

This will need an array of size over 1 million — completely impractical.

Smarter Idea: Prefix Sum

If we look at:


w = [1,2,3,4]
prefixSum = [1, 3, 6, 10]

This prefix sum gives us “boundaries” for cumulative weights.

Imagine breaking the number line [1,10] into:


[1–1] → index 0
[2–3] → index 1
[4–6] → index 2
[7–10] → index 3

Now:


Random number:   Index:
1                0
2,3              1
4,5,6            2
7,8,9,10         3

So again, we have:

Index 0 → 1/10 = 10%
Index 1 → 2/10 = 20%
Index 2 → 3/10 = 30%
Index 3 → 4/10 = 40%

💡 How to Implement

Step 1: Build Prefix Sum Array


prefixSum = [1, 3, 6, 10]

Step 2: Generate a Random Number in Range [1, 10]


int num = rand() % total + 1;

Step 3: Find the Index Such That prefix[i] ≥ num

You can use either:

Linear Search (simple)
Binary Search (faster)
std::lower_bound() in C++ (cleanest)

Binary Search C++ Code


int pickIndex() {
    int num = rand() % prefix.back() + 1;
    int low = 0, high = prefix.size() - 1;
    while (low < high) {
        int mid = low + (high - low) / 2;
        if (num <= prefix[mid]) {
            high = mid;
        } else {
            low = mid + 1;
        }
    }
    return low;
}

Optional C++ One-Liner


return lower_bound(prefix.begin(), prefix.end(), num) - prefix.begin();

Time and Space Complexity

Approach	Time (pick)	Space	Remarks
Naive “Expanded Pool”	O(1)	O(sum(w))	Good for intuition
Prefix + Binary Search	O(log n)	O(n)	Efficient, scalable
Prefix + Linear Search	O(n)	O(n)	Simple, but slower

✅ Summary

Use prefix sums to compress the “repeated pool” idea.
Pick a random number in [1, sum(w)].
Use binary search to find which bucket it falls in.
Return the corresponding index.

This guarantees that pickIndex() will return each index i with probability w[i] / sum(w).

Jdecodes