#528 Random Pick with Weight
the problem
In plain words
You get positive weights w[0..n−1]. Build a class whose pickIndex() returns a random index, where index i comes up with probability w[i] / (sum of all weights). The judge checks that every pick is a valid index and that, over many picks, each index's share is close to its weight's share.
in w = [1]; pickIndex()
out 0
There is only one index.
in w = [1, 3]; pickIndex() × 5
out e.g. [1, 1, 1, 1, 0]
Index 1 should come up about 75% of the time; any sequence with that tendency is accepted.
- 1 ≤ w.length ≤ 10⁴
- 1 ≤ w[i] ≤ 10⁵
- pickIndex is called at most 10⁴ times
your turn
Try it before you peek
Read the examples, say the brute force out loud, then write code: here or on LeetCode ↗. Stuck for 20 minutes? Open one hint at a time.
Solve it in the browser
Design: implement every operation. Compiled by the real javac (Java 17) and judged right here — nothing is sent anywhere.
Solutions are hidden so you can try first. Even 20 minutes of real effort makes the answer stick far better than reading it.
Marking the problem solved, or opening every hint, also shows it.
the aha
Key insight
patternPrefix sumsBinary search (exact & bounds)
where this runsNGINX's random load-balancing method picks an upstream server at random "taking into account weights of servers" — exactly this problem.
Hand out tickets: index i gets w[i] consecutive ticket numbers. Prefix sums tell you where each index's range ends, so draw a uniform ticket and binary-search the first prefix ≥ ticket.
solutions
From brute force to optimal
Fill an array with index i repeated w[i] times and pick a uniform slot.
- total = sum of w; tickets = new int[total].
- Write i into w[i] consecutive slots.
- pickIndex: tickets[random.nextInt(total)].
S = sum of the weights; every slot is written once.
S can reach 10⁴ × 10⁵ = 10⁹ ints, about 4 GB — out of memory at LeetCode's limits.
class Solution { private final int[] tickets; // index i is written w[i] times private final Random random = new Random(); public Solution(int[] w) { int total = 0; for (int x : w) total += x; tickets = new int[total]; // up to 10^4 * 10^5 = 10^9 ints (4 GB) at LeetCode's limits int k = 0; for (int i = 0; i < w.length; i++) for (int c = 0; c < w[i]; c++) tickets[k++] = i; } public int pickIndex() { return indexFor(random.nextInt(tickets.length) + 1); } int indexFor(int ticket) { // ticket is 1..total return tickets[ticket - 1]; }}Store only where each index's tickets end; walk until you pass the ticket.
- prefix[i] = w[0] + … + w[i].
- ticket = random.nextInt(total) + 1.
- Return the first i with prefix[i] ≥ ticket (scan from 0).
The scan may walk the whole array: 10⁴ × 10⁴ = 10⁸ steps for all picks.
The prefix array.
class Solution { private final int[] prefix; private final Random random = new Random(); 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() { return indexFor(random.nextInt(prefix[prefix.length - 1]) + 1); } int indexFor(int ticket) { int i = 0; while (prefix[i] < ticket) i++; // O(n) per pick return i; }}prefix[] is strictly increasing, so the owner of a ticket is a lower bound: binary search it.
- Build prefix[] once.
- ticket = random.nextInt(total) + 1 (1..total, all equally likely).
- lo = 0, hi = n−1; while lo < hi: if prefix[mid] < ticket, lo = mid + 1, else hi = mid.
- Return lo.
14 probes per pick for n = 10⁴.
The prefix array.
class Solution { private final int[] prefix; // index i owns tickets prefix[i-1]+1 .. prefix[i] private final Random random = new Random(); public Solution(int[] w) { prefix = new int[w.length]; int sum = 0; // at most 10^4 * 10^5 = 10^9: fits in an int for (int i = 0; i < w.length; i++) { sum += w[i]; prefix[i] = sum; } } public int pickIndex() { int ticket = random.nextInt(prefix[prefix.length - 1]) + 1; // 1..total, all equally likely return indexFor(ticket); } int indexFor(int ticket) { // the first i with prefix[i] >= ticket int lo = 0, hi = prefix.length - 1; while (lo < hi) { int mid = lo + (hi - lo) / 2; if (prefix[mid] < ticket) lo = mid + 1; else hi = mid; } return lo; }}Dry run
Let the JDK do the search; fix up its negative "not found" result.
- i = Arrays.binarySearch(prefix, ticket).
- Not found → i = −(insertion point) − 1, so return −i − 1.
Same binary search, written by the JDK.
The prefix array.
class Solution { private final int[] prefix; private final Random random = new Random(); 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() { return indexFor(random.nextInt(prefix[prefix.length - 1]) + 1); } int indexFor(int ticket) { int i = Arrays.binarySearch(prefix, ticket); // prefix is strictly increasing because every w[i] >= 1 return i >= 0 ? i : -i - 1; // not found: binarySearch returns -(insertion point) - 1 }}in the room
How to say it in the interview
talk track
- Explain the ticket picture with w = [1, 3] before any code: tickets 1 | 2 3 4.
- Be explicit about the range: nextInt(total) gives 0..total−1, so add 1 and search for prefix ≥ ticket (or skip the +1 and search for prefix > ticket). Off-by-one here is the classic bug.
- State the costs: O(n) build, O(log n) per pick, O(n) memory.
traps
- Off by one in the ticket range: index 0 loses its ticket, or the search runs past the last index.
(int) (Math.random() * total)is fine, butrandom.nextInt(total - 1)silently shrinks the last index's chance.- Building the expanded ticket array: up to 4 GB at LeetCode's limits.
what they ask next
Follow-ups
What if weights change between picks?
Keep the weights in a Fenwick tree: update in O(log n), and find a ticket's owner by walking down the tree in O(log n).
Can a pick be O(1)?
Yes: Walker's alias method (1974) — O(n) setup, then one random index and one coin flip per pick.
before you submit
Edge cases to test
- One weight → always 0
- A tiny weight between huge ones ([100000, 1, 100000]) must still be picked now and then
- The total can reach 10⁹: it still fits in an int, but don't multiply it
practise next
Similar problems
- #303Range Sum Query - ImmutableEasyin Hashing
- #35Search Insert PositionEasyin Binary Search