Random Pick with Weight

medium math binary search prefix sum randomized

Problem

Given an array w where w[i] is the weight of index i, implement pickIndex() so that i is returned with probability w[i] / Σw.

Inputw = [1, 3], target = 2

Output1

Prefix = [1, 4]; first prefix ≥ 2 is at index 1.

w (positive ints, comma-separated)

target

import random
from bisect import bisect_left

class Solution:
    def __init__(self, w):
        self.prefix = []
        total = 0
        for x in w:
            total += x
            self.prefix.append(total)
        self.total = total

    def pickIndex(self):
        t = random.randint(1, self.total)
        return bisect_left(self.prefix, t)

class Solution {
  constructor(w) {
    this.prefix = [];
    let total = 0;
    for (const x of w) { total += x; this.prefix.push(total); }
    this.total = total;
  }
  pickIndex() {
    const t = Math.floor(Math.random() * this.total) + 1;
    let lo = 0, hi = this.prefix.length - 1;
    while (lo < hi) {
      const mid = (lo + hi) >> 1;
      if (this.prefix[mid] < t) lo = mid + 1;
      else hi = mid;
    }
    return lo;
  }
}

class Solution {
    int[] prefix;
    int total;
    Random rand = new Random();
    public Solution(int[] w) {
        prefix = new int[w.length];
        int s = 0;
        for (int i = 0; i < w.length; i++) { s += w[i]; prefix[i] = s; }
        total = s;
    }
    public int pickIndex() {
        int t = rand.nextInt(total) + 1;
        int lo = 0, hi = prefix.length - 1;
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (prefix[mid] < t) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
}

class Solution {
    vector<int> prefix;
    int total;
public:
    Solution(vector<int>& w) {
        int s = 0;
        for (int x : w) { s += x; prefix.push_back(s); }
        total = s;
    }
    int pickIndex() {
        int t = rand() % total + 1;
        return lower_bound(prefix.begin(), prefix.end(), t) - prefix.begin();
    }
};

Explanation

We want index i to come up with probability proportional to its weight w[i]. Picture laying the weights end to end on a number line: index 0 takes a chunk of length w[0], index 1 the next chunk, and so on. Throwing a dart uniformly along that line lands in a bigger chunk more often — exactly the weighting we want.

To do this we first build a prefix-sum array, where prefix[i] is the running total of weights up to index i. This array marks the right edge of each chunk and is always increasing.

For a pick, we sample a random target t in [1, total] (the dart) and find the first chunk whose right edge is at least t. Since prefix is sorted, we use binary search (bisect_left / lower_bound) to find that index in O(log n).

Example: w = [1, 3] gives prefix = [1, 4] and total 4. If t = 2, the first prefix value ≥ 2 is 4 at index 1, so we return 1. Index 1 wins for t in {2,3,4} — that is 3 out of 4 darts, matching its weight 3.

Building the prefix array is a one-time O(n) cost; every pick after that is just a fast binary search.

Time: O(n) init, O(log n) pick Space: O(n)