Log inSign up
mediumLeetCodeNot Started

Random Pick with Weight

A curated StudyFlow practice problem with comprehensive editorial and multi-language solutions.

Arrays
Solve on LeetCode
Lesson

Problem Statement

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith 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) = 1/4 (i.e., 25%), and the probability of picking index 1 is 3 / (1 + 3) = 3/4 (i.e., 75%).

Example 1 · Two Weights:

Input: w = [1, 3]
Output: [0 or 1, with 75% chance for 1]
Explanation: 1 has weight 3 (75% probability), 0 has weight 1 (25% probability).

Example 2 · Three Weights:

Input: w = [1, 2, 3]
Output: [0, 1, or 2 with probs 1/6, 2/6, 3/6]
Explanation: Prefix sums: [1, 3, 6]. Range [1..1] -> 0, [2..3] -> 1, [4..6] -> 2.

Constraints:

  • 1 <= w.length <= 104
  • 1 <= w[i] <= 105
  • pickIndex will be called at most 104 times.

Prerequisites

Before attempting this problem, you should be comfortable with:

  • Prefix Sum & Cumulative Distribution Functions (CDF)
  • Binary Search / Upper Bound lower bound mechanics
  • Pseudo-random number generation in [1, K]

What You'll Learn

Key principles and algorithmic techniques you will master from this problem:

  • Map discrete probability distributions to contiguous intervals via prefix sums.
  • Apply binary search on monotonically increasing cumulative distribution functions (CDF).
  • Generate uniform random targets in 1-indexed range [1, totalSum].

Pattern Recognition Trigger

How to instantly recognize this algorithmic pattern when encountering unseen interview questions:

Signals to Look For:

  • Input structure involves an array or contiguous sequence with constraints on element relationships.
  • Optimal solution requires avoiding O(N2) exhaustive searches through state tracking or pointers.
  • Subproblem exhibits optimal substructure or monotonicity enabling linear or logarithmic processing.

Algorithmic Decision Tree:

  • Build prefix sum array prefix[i] = prefix[i-1] + w[i].
  • Pick random integer in [1, totalSum].
  • Binary search for leftmost index where prefix[idx] >= target.

Connected Patterns:

Alias Method (O(1) sampling)Roulette Wheel Selection in Genetic Algorithms

1. Expanded Frequency Array (Brute Force)

Intuition

If w = [1, 3], create an array of size 1 + 3 = 4: [0, 1, 1, 1]. Picking uniformly at random from this expanded array yields index 0 with probability 1/4 and index 1 with probability 3/4. This is memory-prohibitive when weights are large.

When to apply: Use for tiny weights sum to understand the exact discrete probability distribution.

Expanded Roulette Array (O(Sum(W)) Space)

2 of 2 questions left today
Step 1 / 6
Trace timeline17% explored
Current operationSTEP 1 OF 6
Direct expansion models probability through duplication
Visual state
Values and pointers at this exact step
pointers
S
idx: 0
idx: 1
idx: 1
idx: 1
0
↑ roulette
1
2
3
Variables at this step
totalWeight
4
array
[0,1,1,1]
Check / comparison
Expand w = [1, 3] into roulette array: one '0' and three '1's
Decision / next transition
Total length = 1 + 3 = 4.
Why this step matters

Direct expansion models probability through duplication.

Step 1 / 6
JAVA SOLUTIONCLEAN & OPTIMIZED
1import java.util.*;
2
3public class Main {
4 public static class Solution {
5 private List<Integer> list = new ArrayList<>();
6 private Random rand = new Random();
7
8 public Solution(int[] w) {
9 for (int i = 0; i < w.length; i++) {
10 for (int count = 0; count < w[i]; count++) {
11 list.add(i);
12 }
13 }
14 }
15
16 public int pickIndex() {
17 return list.get(rand.nextInt(list.size()));
18 }
19 }
20
21 public static void main(String[] args) {
22 int[] w = {1, 3};
23 Solution sol = new Solution(w);
24 System.out.println("Random Pick: " + sol.pickIndex());
25 }
26}

Time & Space Complexity

  • Time complexity: O(sum(w)) space, O(1) pick — Building the expanded array takes time proportional to the sum of weights.
  • Space complexity: O(sum(w)) — Stores sum(w) entries, which exceeds memory when weights sum to 109.

Why Brute Force Doesn't Scale:

For N = 105, quadratic O(N2) approaches require ~1010 operations, exceeding the standard 1-second CPU budget.

The optimal approach processes N elements in a single or logarithmic pass, completing in under 10ms with minimal memory footprint.

2. Prefix Sums + Linear Search (Better)

Intuition

Compute prefix sums of weights: prefix[i] = sum(w[0..i]). The random target is a uniformly distributed integer in [1, prefix[n-1]]. A linear scan from index 0 finds the first index i where prefix[i] >= target.

When to apply: Use to reduce memory to O(n) before optimizing query time with binary search.

Prefix Sum Linear Scan (O(N) Time per Pick)

2 of 2 questions left today
Step 1 / 6
Trace timeline17% explored
Current operationSTEP 1 OF 6
Prefix sums reduce memory from O(Sum(W)) to O(N)
Visual state
Values and pointers at this exact step
pointers
S
1
4
0
↑ pref
1
Variables at this step
prefix
[1,4]
totalWeight
4
Check / comparison
Compute prefix sums: prefix[0] = 1, prefix[1] = 1 + 3 = 4
Decision / next transition
Total weight = 4. Intervals: index 0 covers [1..1], index 1 covers [2..4].
Why this step matters

Prefix sums reduce memory from O(Sum(W)) to O(N).

Step 1 / 6
JAVA SOLUTIONCLEAN & OPTIMIZED
1import java.util.*;
2
3public class Main {
4 public static class Solution {
5 private int[] prefix;
6 private int total;
7 private Random rand = new Random();
8
9 public Solution(int[] w) {
10 prefix = new int[w.length];
11 prefix[0] = w[0];
12 for (int i = 1; i < w.length; i++) {
13 prefix[i] = prefix[i - 1] + w[i];
14 }
15 total = prefix[w.length - 1];
16 }
17
18 public int pickIndex() {
19 int target = rand.nextInt(total) + 1;
20 for (int i = 0; i < prefix.length; i++) {
21 if (prefix[i] >= target) return i;
22 }
23 return prefix.length - 1;
24 }
25 }
26
27 public static void main(String[] args) {
28 int[] w = {1, 3, 2, 4};
29 Solution sol = new Solution(w);
30 System.out.println("Random Pick: " + sol.pickIndex());
31 }
32}

Time & Space Complexity

  • Time complexity: O(n) init, O(n) pick — Building prefix sums takes O(n). Linear scan takes up to O(n) per pick.
  • Space complexity: O(n) — O(n) memory to store the prefix sum array.

3. Prefix Sums + Binary Search (Optimal)

Intuition

Since weights are positive, the prefix sum array is strictly increasing and sorted. We sample target in [1, totalSum]. Binary search finds the smallest index mid such that prefix[mid] >= target in O(log n) time per call with O(n) preprocessing.

When to apply: Use as the optimal production solution for repeated random weight selections.

Prefix Sum + Binary Search (O(log N) Time per Pick)

2 of 2 questions left today
Step 1 / 6
Trace timeline17% explored
Current operationSTEP 1 OF 6
Strictly increasing prefix sums guarantee binary search correctness
Visual state
Values and pointers at this exact step
pointers
S
1
4
0
↑ prefix
1
Variables at this step
prefix
[1,4]
total
4
Check / comparison
Precompute cumulative prefix sums: prefix = [1, 4]. Total sum = 4
Decision / next transition
Because weights > 0, prefix array is strictly increasing.
Why this step matters

Strictly increasing prefix sums guarantee binary search correctness.

Step 1 / 6
JAVA SOLUTIONCLEAN & OPTIMIZED
1import java.util.*;
2
3public class Main {
4 public static class Solution {
5 private int[] prefix;
6 private int total;
7 private Random rand = new Random();
8
9 public Solution(int[] w) {
10 prefix = new int[w.length];
11 prefix[0] = w[0];
12 for (int i = 1; i < w.length; i++) {
13 prefix[i] = prefix[i - 1] + w[i];
14 }
15 total = prefix[w.length - 1];
16 }
17
18 public int pickIndex() {
19 int target = rand.nextInt(total) + 1;
20 int low = 0, high = prefix.length - 1;
21 while (low < high) {
22 int mid = low + (high - low) / 2;
23 if (prefix[mid] >= target) {
24 high = mid;
25 } else {
26 low = mid + 1;
27 }
28 }
29 return low;
30 }
31 }
32
33 public static void main(String[] args) {
34 int[] w = {1, 3, 2, 4};
35 Solution sol = new Solution(w);
36 System.out.println("Random Pick: " + sol.pickIndex());
37 }
38}

Time & Space Complexity

  • Time complexity: O(n) init, O(log n) pick — One O(n) initialization pass to compute prefix sums; each pickIndex call takes O(log n) binary search.
  • Space complexity: O(n) — O(n) space to hold the prefix sum array.

Approach Comparison Matrix

Side-by-side comparison of time, space, and algorithmic trade-offs across each approach:

ApproachTimeSpaceCore IdeaTrade-off / Bottleneck
1.Brute Force — Expanded Bucket ArrayO(sum(w)) init, O(1) pickO(sum(w))Fill an array with w[i] copies of index i, then pick uniformly at random.Massive memory footprint.
2.Better — Linear Scan on Prefix SumsO(N) init, O(N) pickO(N)Generate target in [1, totalSum] and scan prefix sums sequentially.Sequential search on prefix sums.
3.Optimal — Binary Search on Prefix SumsOptimalO(N) init, O(log N) pickO(N)Binary search on monotonic prefix sums to find the bucket index in O(log N).Optimal for cumulative frequency search.

Common Pitfalls & Crucial Edge Cases

Off-by-One Errors

Incorrect boundary indexing when adjusting pointers or array slices.

Integer Overflow

Using 32-bit signed integer types when accumulative operations exceed 231 - 1 in C++/Java.

Unnecessary Memory Allocation

Allocating large auxiliary arrays or objects inside inner loops causing GC overhead.

Failing to Reset Loop State

Not re-initializing state variables or counters between iterations or subproblems.

Crucial Edge Cases to Validate:

  • Empty or Single Element Array — Handle minimal input sizes (N = 0 or N = 1) cleanly without out-of-bounds errors.
  • Duplicate Elements — Verify behavior when all elements are identical (e.g. all 0s or all equal values).
  • Extreme Numerical Values — Account for maximum/minimum allowed constraint values (-109 to 109) to prevent overflow.
  • Boundary Conditions — Ensure correct execution when target or match is located at the first or last index.

Interview Tips & Best Practices

How to communicate your thought process and address follow-ups in a live technical interview:

  1. Always state the brute force baseline before presenting the optimal solution to demonstrate structured problem solving.
  2. Explicitly verify boundary conditions (empty arrays, single elements, duplicates, extremes) out loud with the interviewer.
  3. Analyze both time and auxiliary space complexities clearly, explaining trade-offs between speed and memory allocation.

Quick Comprehension Check

Test your intuitive understanding of the core concept in 30 seconds:

1.What is the optimal time and auxiliary space complexity for Random Pick with Weight?

2.Why is the brute force approach suboptimal for large inputs in Random Pick with Weight?

3.Which data structure or technique enables the optimal solution for Random Pick with Weight?

Back to DSA Sheet