Skip to main content
CodePeet

Random Pick with Weight

Try it yourself
MediumProblemExternal link

Description

You get an array of positive integers w, where w[i] is the weight of index i. Build a class that picks an index at random, so that index i comes up with probability w[i] / sum(w).

A judge can't check random answers exactly, so here the random number is drawn for you and passed in. Line up the whole numbers from 1 to sum(w), and give index 0 the first w[0] of them, index 1 the next w[1], and so on. pickIndex(r) receives one of those numbers, r, drawn uniformly at random, and returns the index that owns it. Index i owns w[i] of the sum(w) numbers, so it's picked with probability exactly w[i] / sum(w).

Implement the class:

  • Solution(w) builds it from the weights.
  • pickIndex(r) returns the index that owns the draw r.

Each test gives w and a list of draws. The class is built once, pickIndex is called once per draw, and the output is the list of picked indexes.

Examples

Example 1

Input: w = [1], draws = [1]

Output: [0]

Explanation: There's only index 0, so every draw picks it.

Example 2

Input: w = [1, 3], draws = [1, 2, 3, 4, 2]

Output: [0,1,1,1,1]

Explanation: Index 0 owns the draw 1 and index 1 owns 2, 3 and 4, so a uniform draw picks index 1 three times as often as index 0: 3 out of 4.

Example 3

Input: w = [3, 1, 2], draws = [3, 4, 5, 1, 6]

Output: [0,1,2,0,2]

Explanation: Index 0 owns 1 to 3, index 1 owns 4, and index 2 owns 5 and 6. The draw 3 is the last one index 0 owns, and 4 is the first (and only) one of index 1.

Constraints

  • 1 ≤ w.length ≤ 10^4
  • 1 ≤ w[i] ≤ 10^5
  • 1 ≤ number of draws ≤ 10^4
  • 1 ≤ r ≤ sum(w)

Editorial

Brute Force

Intuition

Walk the indexes in order and keep a running total of the weights. After adding w[i], the running total is the last number index i owns. So the first index whose running total reaches r is the one that owns r.

Step-by-Step Explanation

Let's find the owner of r = 13 in w = [2, 5, 1, 4, 3, 6, 2, 1]:

  1. Add w[0] = 2, so the running total is 2 and index 0 owns 1 to 2. 2 is still below 13, so keep going.
  2. Add w[1] = 5, so the running total is 7 and index 1 owns 3 to 7. 7 is still below 13, so keep going.
  3. Add w[2] = 1, so the running total is 8 and index 2 owns 8 to 8. 8 is still below 13, so keep going.
  4. Add w[3] = 4, so the running total is 12 and index 3 owns 9 to 12. 12 is still below 13, so keep going.
  5. Add w[4] = 3, so the running total is 15 and index 4 owns 13 to 15. 13 is in there, so the answer is 4.

That's index 4, after 5 steps.

Walk the Weights — Each step adds the next weight to the running total. Amber means the total is still below r; green is the first index whose total reaches it.

Algorithm

  1. In the constructor, keep w.
  2. In pickIndex(r): set running = 0; for each index i, add w[i] to running, and return i as soon as running ≥ r.

Code

class Solution {
    vector<int> w;

public:
    Solution(vector<int>& w) : w(w) {}

    int pickIndex(int r) {
        int running = 0;
        for (int i = 0; i < (int)w.size(); i++) {
            running += w[i];                    // index i owns the numbers up to running
            if (running >= r) return i;
        }
        return (int)w.size() - 1;               // not reached: r is at most sum(w)
    }
};

class Solution:
    def __init__(self, w: list[int]):
        self.w = w

    def pickIndex(self, r: int) -> int:
        running = 0
        for i, weight in enumerate(self.w):
            running += weight                   # index i owns the numbers up to running
            if running >= r:
                return i
        return len(self.w) - 1                  # not reached: r is at most sum(w)

class Solution {
    private final int[] w;

    public Solution(int[] w) {
        this.w = w;
    }

    public int pickIndex(int r) {
        int running = 0;
        for (int i = 0; i < w.length; i++) {
            running += w[i];                    // index i owns the numbers up to running
            if (running >= r) return i;
        }
        return w.length - 1;                    // not reached: r is at most sum(w)
    }
}

var Solution = function(w) {
    this.w = w;
};

Solution.prototype.pickIndex = function(r) {
    let running = 0;
    for (let i = 0; i < this.w.length; i++) {
        running += this.w[i];                   // index i owns the numbers up to running
        if (running >= r) return i;
    }
    return this.w.length - 1;                   // not reached: r is at most sum(w)
};

Complexity Analysis

Time Complexity: O(n) per pick

A draw in the last index's block walks every weight.

Space Complexity: O(1) besides w

Why This Approach Is Not Efficient

With 10^4 weights and 10^4 draws that all land in the last block, every pick walks all 10^4 weights: 10^8 additions in all, which is over the time limit in Python.

And every pick recomputes the same running totals: they don't depend on r. Compute them once in the constructor instead. Every weight is positive, so the totals only grow, and a sorted array is something binary search can search.

Optimal Approach - Running Totals and Binary Search

Intuition

Build totals[i] = w[0] + … + w[i] once. Index i owns the numbers after totals[i − 1] up to totals[i], so the owner of r is the first index whose total is at least r. That's a first yes on a sorted array, the lower bound from Introduction to Binary Search.

Mind the ≥. A draw equal to totals[i] is the last number index i owns, so it belongs to i, not to i + 1. And the last total is sum(w), which every draw is at most, so the answer is always inside the array: start with hi at the last index.

Step-by-Step Explanation

Let's find the owner of r = 13 again. The running totals of [2, 5, 1, 4, 3, 6, 2, 1] are [2, 7, 8, 12, 15, 21, 23, 24]:

  1. totals[3] = 12 < 13, a no, so lo becomes 4.
  2. totals[5] = 21 ≥ 13, a yes, so hi becomes 5.
  3. totals[4] = 15 ≥ 13, a yes, so hi becomes 4.

Index 4 again, in 3 looks instead of 5 steps.

Search the Running Totals — The row is the running totals. Each step looks at the middle one: green when it reaches r, amber when it doesn't.

Algorithm

  1. In the constructor, build totals: totals[i] = totals[i − 1] + w[i].
  2. In pickIndex(r): set lo = 0, hi = n − 1. While lo < hi, take mid; if totals[mid] ≥ r, set hi = mid, otherwise lo = mid + 1.
  3. Return lo.

Code

class Solution {
    vector<int> totals;                         // totals[i] = w[0] + ... + w[i]: the last draw index i owns

public:
    Solution(vector<int>& w) {
        int running = 0;                        // at most 10^4 * 10^5 = 10^9: fits an int
        for (int weight : w) {
            running += weight;
            totals.push_back(running);
        }
    }

    int pickIndex(int r) {
        int lo = 0, hi = (int)totals.size() - 1;    // the first i with totals[i] >= r; the last is sum(w) >= r
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (totals[mid] >= r) hi = mid;     // yes: index mid owns r or a draw after it
            else lo = mid + 1;                  // no: r lies past index mid's block
        }
        return lo;
    }
};

class Solution:
    def __init__(self, w: list[int]):
        self.totals = []                        # totals[i] = w[0] + ... + w[i]: the last draw index i owns
        running = 0
        for weight in w:
            running += weight
            self.totals.append(running)

    def pickIndex(self, r: int) -> int:
        lo, hi = 0, len(self.totals) - 1        # the first i with totals[i] >= r; the last total is sum(w) >= r
        while lo < hi:
            mid = (lo + hi) // 2
            if self.totals[mid] >= r:
                hi = mid                        # yes: index mid owns r or a draw after it
            else:
                lo = mid + 1                    # no: r lies past index mid's block
        return lo

class Solution {
    private final int[] totals;                 // totals[i] = w[0] + ... + w[i]: the last draw index i owns

    public Solution(int[] w) {
        totals = new int[w.length];
        int running = 0;                        // at most 10^4 * 10^5 = 10^9: fits an int
        for (int i = 0; i < w.length; i++) {
            running += w[i];
            totals[i] = running;
        }
    }

    public int pickIndex(int r) {
        int lo = 0, hi = totals.length - 1;     // the first i with totals[i] >= r; the last total is sum(w) >= r
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (totals[mid] >= r) hi = mid;     // yes: index mid owns r or a draw after it
            else lo = mid + 1;                  // no: r lies past index mid's block
        }
        return lo;
    }
}

var Solution = function(w) {
    this.totals = [];                           // totals[i] = w[0] + ... + w[i]: the last draw index i owns
    let running = 0;
    for (const weight of w) {
        running += weight;
        this.totals.push(running);
    }
};

Solution.prototype.pickIndex = function(r) {
    let lo = 0, hi = this.totals.length - 1;    // the first i with totals[i] >= r; the last total is sum(w) >= r
    while (lo < hi) {
        const mid = lo + Math.floor((hi - lo) / 2);
        if (this.totals[mid] >= r) hi = mid;    // yes: index mid owns r or a draw after it
        else lo = mid + 1;                      // no: r lies past index mid's block
    }
    return lo;
};

Complexity Analysis

Time Complexity: O(n) to build, O(log n) per pick

About 14 looks for 10^4 weights.

Space Complexity: O(n) for the totals

The largest total is 10^4 × 10^5 = 10^9, which fits in a 32-bit int.

Drawing r Yourself

On LeetCode, pickIndex() takes no argument: it draws r itself, uniformly from 1 to sum(w), and then runs the same search. In Python that's random.randint(1, total), in Java ThreadLocalRandom.current().nextInt(1, total + 1), in JavaScript Math.floor(Math.random() * total) + 1, and in C++ a uniform_int_distribution<int>(1, total) over an mt19937. Avoid plain rand() in C++: it can top out at 32,767 on some compilers, far below a total of 10^9, so most indexes would never come up.

Does it really pick by weight? 100,000 seeded draws on w = [1, 3], run on the judge, picked index 0 25,044 times and index 1 74,956 times: a quarter and three quarters.

Common Mistakes

Each of these looks right on a quick read, and the tests still reject it. The outputs below come from running the code.

Looking for the first total above r. if self.totals[mid] > r:

On w = [1, 3] with the draws [1, 2, 3, 4, 2], it returns [1,1,1,1,1]; the answer is [0,1,1,1,1]. totals[i] is the LAST draw index i owns, so r = totals[i] still belongs to index i. Asking for the first total greater than r skips to index i + 1 whenever r lands exactly on the end of a block: with w = [1, 3], the draw 1 picks index 1 instead of 0.

Keeping a 0 in front of the totals. self.totals = [0] # then the running totals

On w = [3, 1, 2] with the draws [3, 4, 5, 1, 6], it returns [1,2,3,1,3]; the answer is [0,1,2,0,2]. Starting the totals with 0 shifts every position by one: the first total that reaches r is then at index i + 1 for the block that owns r. Either drop the leading 0 or return the position minus 1.

Searching the weights instead of the running totals. self.totals.append(weight)

On w = [3, 1, 2] with the draws [3, 4, 5, 1, 6], it returns [2,2,2,0,2]; the answer is [0,1,2,0,2]. The weights aren't sorted, and a single weight says nothing about where its block ends. Binary search needs the running totals: they only grow, and totals[i] is where index i's block ends.

Looping while lo ≤ hi with hi = mid. while lo <= hi:

On w = [1] with the draws [1], it never finishes: the judge stops it at the time limit; the answer is [0]. When lo and hi meet, mid is lo and the update hi = mid changes nothing, so the loop never ends. The first-yes loop runs while lo < hi and stops with lo == hi on the answer.

Where This Pattern Is Taught

This problem practises the lower bound (the first value at least r in a sorted array), on running totals. It's taught in Introduction to Binary Search.