Random Pick with Weight - Solution & Explanation

MediumArray Math Binary Search Prefix Sum15 min readAsked at: Amazon, Microsoft, Apple +23

Problem Statement

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%), and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).

Example 1:

Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]

Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.

Example 2:

Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]

Explanation
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // return 1. It is returning the second element (index = 1) that has a probability of 3/4.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It is returning the first element (index = 0) that has a probability of 1/4.

Since this is a randomization problem, multiple answers are allowed.
All of the following outputs can be considered correct:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
and so on.

Constraints:

1 <= w.length <= 10⁴
1 <= w[i] <= 10⁵
pickIndex will be called at most 10⁴ times.

Approach Overview

Problem Overview: You receive an array w where w[i] represents the weight of index i. The task is to implement pickIndex() so that the probability of picking index i is proportional to its weight. If w = [1,3], index 1 should be returned three times more often than index 0.

The core idea is to treat weights as ranges on a number line. If the total weight is 10, you randomly choose a number from [1,10] and determine which weight range contains that number.

Approach 1: Cumulative Sum with Linear Search (Initialization: O(n), Pick: O(n) time, O(n) space)

Build a prefix sum array where prefix[i] stores the total weight from index 0 to i. For example, weights [1,3,2] become prefix sums [1,4,6]. During pickIndex(), generate a random integer between 1 and the total weight. Iterate through the prefix array and return the first index where prefix[i] >= target. This works because each index occupies a proportional segment of the total range. The implementation is simple and demonstrates the key prefix sum idea, but every pick requires scanning the array.

Approach 2: Cumulative Sum with Binary Search (Initialization: O(n), Pick: O(log n) time, O(n) space)

The optimized approach keeps the same prefix sum array but replaces the linear scan with binary search. After generating a random number between 1 and the total weight, perform a lower-bound search on the prefix array to find the first index where the prefix value is greater than or equal to the target. Because the prefix array is sorted by construction, binary search finds the correct index in O(log n). This significantly improves performance when pickIndex() is called many times. The method combines array preprocessing, prefix sums, and randomized selection.

Recommended for interviews: The prefix sum with binary search solution is the expected answer. It shows that you recognize the weighted probability mapping and know how to convert it into a range search problem. Mentioning the linear scan approach first demonstrates understanding of the probability model, while the binary search optimization shows algorithmic maturity and awareness of repeated-query performance.

Approach 1: Cumulative Sum with Linear Search

This approach involves creating a cumulative sum array based on the provided weights. The idea is to convert the weight array into a cumulative distribution, where each element represents the summed result of all previous weights including the current one. When we generate a random number, we search for its position in this cumulative array to determine which index to return.

This C solution implements a cumulative sum array ('prefixSum') to store the cumulative weights of the input array 'w'. The 'solutionPickIndex' function then generates a random number within the range of the total sum of weights and searches for the corresponding index in the cumulative sum array. The search is done linearly here.

Code

C C++Java Python C#JavaScript

C++

Java

Python

JavaScript

Complexity

Time Complexity: O(N) for initialization, O(N) for the pickIndex.
Space Complexity: O(N) due to the cumulative sum storage.

Try this approach in the editor →

Approach 2: Cumulative Sum with Binary Search

This optimized approach also uses a cumulative sum array, but instead of performing a linear search to find the appropriate index, we use a binary search. This greatly improves the efficiency when determining which index corresponds to a given cumulative value, especially beneficial for larger arrays.

This C implementation uses binary search to efficiently find the index corresponding to the random target. After setting up the cumulative sum array, a target number is generated, and binary search is used to pinpoint the index.

Code

C C++Java Python C#JavaScript

C++

Java

Python

JavaScript

Complexity

Time Complexity: O(N) for initialization, O(log N) for pickIndex.
Space Complexity: O(N) for the cumulative sum array.

Try this approach in the editor →

Approach 3: Default Approach

Code

Python Java C++Go Rust JavaScript

Python

Java

C++

Rust

JavaScript

Try this approach in the editor →

Complexity Comparison

Approach	Complexity
Cumulative Sum with Linear Search	Time Complexity: O(N) for initialization, O(N) for the pickIndex. Space Complexity: O(N) due to the cumulative sum storage.
Cumulative Sum with Binary Search	Time Complexity: O(N) for initialization, O(log N) for pickIndex. Space Complexity: O(N) for the cumulative sum array.
Default Approach	—

Detailed Complexity Analysis

Approach	Time	Space	When to Use
Cumulative Sum with Linear Search	Init: O(n), Pick: O(n)	O(n)	Simple implementation or when pickIndex is called only a few times
Cumulative Sum with Binary Search	Init: O(n), Pick: O(log n)	O(n)	Best choice when many random picks are required

Video Solution

Random Pick with Weight | Bucket approach | Leetcode #528 | Binary search • Techdose • 42,813 views views

Watch 9 more video solutions →

Frequently Asked Questions

Is Random Pick with Weight easy or hard?

Random Pick with Weight is generally classified as a medium difficulty problem. The main challenge is recognizing that weighted probability can be converted into prefix sum ranges and then optimized with binary search for repeated queries.

How to solve Random Pick with Weight in O(n)?

Create a prefix sum array where each element represents the cumulative weight up to that index. Generate a random number between 1 and the total weight, then linearly scan the prefix array until the cumulative value is greater than or equal to the random number. Initialization is O(n) and each pick is O(n).

What is the best approach for Random Pick with Weight?

The most efficient approach uses a prefix sum array combined with binary search. Build cumulative weights during initialization, then generate a random number between 1 and the total weight and locate its position using binary search. This gives O(n) preprocessing time and O(log n) time per pick.

What data structure is used in Random Pick with Weight?

The core data structure is a prefix sum array that stores cumulative weights. This array is paired with a random number generator and binary search to efficiently map a random value to the correct weighted index.

What is the time complexity of Random Pick with Weight?

Building the prefix sum array takes O(n) time and O(n) space. Each call to pickIndex() runs in O(log n) using binary search on the prefix array. A simpler linear scan version exists but takes O(n) per pick.

Random Pick with Weight Python or Java solution approach?

Both Python and Java implementations follow the same pattern: build a prefix sum array in the constructor, compute the total weight, generate a random number within that range, and use binary search to locate the index whose prefix sum covers that value.

Is Random Pick with Weight asked at Google, Amazon, or Meta?

Random Pick with Weight appears frequently in interviews focused on probability and prefix sums. Variants of the problem have been reported in Amazon, Google, and Meta interview preparation lists because it tests randomized algorithms and binary search optimization.

Ready to solve this problem?

Practice Random Pick with Weight with our built-in code editor and test cases.

Practice on FleetCode

Random Pick Index

Random Pick with Blacklist

Problem Info

DifficultyMedium

Acceptance49.0%

Approaches3

Reading time15 min

Asked at

Amazon Microsoft Apple Meta Sony

Practice this problem

Open in Editor

Problem Statement

Approach Overview

Approach 1: Cumulative Sum with Linear Search

Code

Complexity

Approach 2: Cumulative Sum with Binary Search

Code

Complexity

Approach 3: Default Approach

Code

Complexity Comparison

Detailed Complexity Analysis

Video Solution

Frequently Asked Questions

Related Problems

Ready to solve this problem?

Problem Info

Table of Contents