Random Pick with Weight - Practice Coding Problems

Random Pick with Weight - Problem

Weighted Random Selection

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the i^th index. Your task is to implement a data structure that supports weighted random selection.

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 index i should be proportional to its weight: w[i] / sum(w).

Example: If w = [1, 3], then:
• Probability of picking index 0 = 1 / (1 + 3) = 0.25 (25%)
• Probability of picking index 1 = 3 / (1 + 3) = 0.75 (75%)

This means index 1 should be selected 3 times more often than index 0 over many calls to pickIndex().

Input & Output

example_1.py — Basic Usage

$ Input: w = [1, 3] Solution(w) pickIndex() called multiple times

› Output: Returns 0 approximately 25% of the time, 1 approximately 75% of the time

💡 Note: With weights [1, 3], the total sum is 4. Index 0 has weight 1/4 = 25% probability, index 1 has weight 3/4 = 75% probability.

example_2.py — Equal Weights

$ Input: w = [1, 1, 1, 1] Solution(w) pickIndex()

› Output: Each index (0, 1, 2, 3) returned with equal 25% probability

💡 Note: When all weights are equal, each index has the same probability of being selected, making it uniform random selection.

example_3.py — Single Element

$ Input: w = [5] Solution(w) pickIndex()

› Output: Always returns 0

💡 Note: With only one element, there's only one possible index to return, so it's selected with 100% probability.

Visualization

Tap to expand

Understanding the Visualization

Create Weighted Dartboard

Each index becomes a sector with size proportional to its weight

Build Prefix Sum Boundaries

Calculate cumulative boundaries to know where each sector starts and ends

Throw the Dart

Generate a random position on the dartboard

Binary Search for Sector

Quickly find which sector the dart landed in using binary search

Key Takeaway

🎯 Key Insight: Transform the weighted selection problem into a range lookup problem using prefix sums, then leverage binary search for efficient O(log n) queries!

Time & Space Complexity

Time Complexity

⏱️

O(n)

Each pickIndex() call requires O(n) time to iterate through weights array

✓ Linear Growth

Space Complexity

O(1)

Only using constant extra space for variables

✓ Linear Space

Constraints

1 ≤ w.length ≤ 10⁴
1 ≤ w[i] ≤ 10⁵
pickIndex will be called at most 10⁴ times
All weights are positive integers

Asked in

G Google 42 f Facebook 38 a Amazon 35 ⊞ Microsoft 28

The optimal solution uses prefix sums combined with binary search. During initialization, compute cumulative weights in O(n) time. For each query, generate a random number and use binary search to find the corresponding index in O(log n) time. This approach efficiently handles repeated queries while maintaining proper probability distribution.

Common Approaches

Approach	Time	Space	Notes
✓ Linear Search Each Time	O(n)	O(1)	For each pick, iterate through weights and use random selection
Prefix Sum + Binary Search (Optimal)	O(log n)	O(n)	Precompute prefix sums and use binary search for O(log n) queries

Linear Search Each Time — Algorithm Steps

Calculate the total sum of all weights
Generate a random number between 0 and total_sum
Iterate through the weights array
Keep a running sum and return the index when random number falls in current weight's range

Visualization

Tap to expand

Step-by-Step Walkthrough

Generate Random

Generate random number between 0 and total weight sum

Linear Search

Iterate through weights, maintaining cumulative sum

Find Range

Return index when random number falls within current weight range

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <time.h>

typedef struct {
    int* weights;
    int size;
} Solution;

Solution* solutionCreate(int* w, int wSize) {
    Solution* obj = (Solution*)malloc(sizeof(Solution));
    obj->weights = (int*)malloc(wSize * sizeof(int));
    obj->size = wSize;
    
    for (int i = 0; i < wSize; i++) {
        obj->weights[i] = w[i];
    }
    
    srand(time(NULL));
    return obj;
}

int solutionPickIndex(Solution* obj) {
    int total = 0;
    for (int i = 0; i < obj->size; i++) {
        total += obj->weights[i];
    }
    
    double target = ((double)rand() / RAND_MAX) * total;
    int cumsum = 0;
    
    for (int i = 0; i < obj->size; i++) {
        cumsum += obj->weights[i];
        if (target <= cumsum) {
            return i;
        }
    }
    return obj->size - 1;
}

void solutionFree(Solution* obj) {
    free(obj->weights);
    free(obj);
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

Each pickIndex() call requires O(n) time to iterate through weights array

✓ Linear Growth

Space Complexity

O(1)

Only using constant extra space for variables

✓ Linear Space

Constraints

1 ≤ w.length ≤ 10⁴
1 ≤ w[i] ≤ 10⁵
pickIndex will be called at most 10⁴ times
All weights are positive integers

67.2K Views

High Frequency

~22 min Avg. Time

1.8K Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Visualization

Time & Space Complexity

Related Problems

Constraints

Common Approaches

Linear Search Each Time — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Constraints

Select Compiler