Random Pick with Weight

Random Pick with Weight - Problem

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%).

Input & Output

Example 1 — Basic Weighted Selection

$ Input: w = [1,3]

› Output: 0 or 1 (with probabilities 25% and 75%)

💡 Note: Index 0 has weight 1, index 1 has weight 3. Total sum is 4. Random number 1 maps to index 0, numbers 2-4 map to index 1.

Example 2 — Equal Weights

$ Input: w = [2,2,2]

› Output: 0, 1, or 2 (each with 33.33% probability)

💡 Note: All indices have equal weight 2. Total sum is 6. Numbers 1-2 → index 0, 3-4 → index 1, 5-6 → index 2.

Example 3 — Single Element

$ Input: w = [5]

› Output: 0 (100% probability)

💡 Note: Only one index available, so it's always selected regardless of weight value.

Constraints

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

Visualization

Tap to expand

Asked in

F Facebook 35 G Google 28 a Amazon 22 M Microsoft 18

The key insight is to convert weights into cumulative probability ranges using prefix sums, then use binary search to efficiently find which range a random number falls into. The optimal approach uses prefix sum + binary search for O(log n) pick time and O(n) space.

Common Approaches

✓ Binary Search on Prefix Sums

⏱️ Time: O(log n) Space: O(n)

Build prefix sum array during initialization, then for each pick operation, generate a random number and use binary search to efficiently find which weight range it falls into.

Linear Search with Prefix Sum

⏱️ Time: O(n) Space: O(n)

Build a prefix sum array during initialization, then for each pick, generate a random number and linearly search through the prefix sums to find which weight range it falls into.

Binary Search on Prefix Sums — Algorithm Steps

Build prefix sum array from weights
Generate random number from 1 to total sum
Binary search on prefix sums to find target range in O(log n)

Visualization

Tap to expand

Step-by-Step Walkthrough

Build Prefix Sums

Convert weights [1,3,2] to cumulative [1,4,6]

Generate Random

Pick random number from 1 to 6 (total sum)

Binary Search

Use binary search to find target in O(log n) time

Code -

solution.c — C

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

typedef struct {
    int* prefixSums;
    int totalSum;
    int size;
} Solution;

Solution* solutionCreate(int* w, int wSize) {
    Solution* obj = (Solution*)malloc(sizeof(Solution));
    obj->prefixSums = (int*)malloc(wSize * sizeof(int));
    obj->size = wSize;
    int prefixSum = 0;
    for (int i = 0; i < wSize; i++) {
        prefixSum += w[i];
        obj->prefixSums[i] = prefixSum;
    }
    obj->totalSum = prefixSum;
    return obj;
}

int solutionPickIndex(Solution* obj) {
    int target = rand() % obj->totalSum + 1;
    int left = 0, right = obj->size - 1;
    while (left < right) {
        int mid = left + (right - left) / 2;
        if (obj->prefixSums[mid] < target) {
            left = mid + 1;
        } else {
            right = mid;
        }
    }
    return left;
}

void solution(int* w, int wSize) {
    if (wSize == 0) {
        printf("0\n");
        return;
    }
    
    // For arrays with only one element, always return "0"
    if (wSize == 1) {
        printf("0\n");
        return;
    }
    
    Solution* obj = solutionCreate(w, wSize);
    
    // Run multiple times to show different possible outcomes
    static int results[1000];
    int resultsCount = 0;
    int maxAttempts = 1000;
    
    for (int i = 0; i < maxAttempts; i++) {
        int pick = solutionPickIndex(obj);
        
        // Check if this result is already in our set
        int found = 0;
        for (int j = 0; j < resultsCount; j++) {
            if (results[j] == pick) {
                found = 1;
                break;
            }
        }
        
        if (!found) {
            results[resultsCount++] = pick;
        }
        
        // If we have all possible results, break early
        if (resultsCount >= wSize) {
            break;
        }
    }
    
    // Sort results
    for (int i = 0; i < resultsCount - 1; i++) {
        for (int j = i + 1; j < resultsCount; j++) {
            if (results[i] > results[j]) {
                int temp = results[i];
                results[i] = results[j];
                results[j] = temp;
            }
        }
    }
    
    // Return all possible results in the expected format
    printf("any of [");
    for (int i = 0; i < resultsCount; i++) {
        if (i > 0) printf(", ");
        printf("%d", results[i]);
    }
    printf("]\n");
    
    free(obj->prefixSums);
    free(obj);
}

int main() {
    srand((unsigned int)(time(NULL) ^ getpid()));
    
    char line[1000];
    if (!fgets(line, sizeof(line), stdin)) {
        printf("0\n");
        return 0;
    }
    
    static int w[100];
    int size = 0;
    
    char* start = strchr(line, '[');
    if (start) {
        start++;
        char* end = strchr(start, ']');
        if (end) {
            *end = '\0';
        }
        
        char* ptr = start;
        while (*ptr) {
            while (*ptr == ' ' || *ptr == ',' || *ptr == '\t' || *ptr == '\n') ptr++;
            if (*ptr == '\0') break;
            
            char* endptr;
            long val = strtol(ptr, &endptr, 10);
            if (ptr != endptr) {
                w[size++] = (int)val;
                ptr = endptr;
            } else {
                break;
            }
        }
    }
    
    solution(w, size);
    
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(log n)

Each pickIndex() call uses binary search on n prefix sums

⚡ Linearithmic

Space Complexity

O(n)

Store prefix sum array of size n

⚡ Linearithmic Space

156.0K Views

Medium Frequency

~25 min Avg. Time

1.9K 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

Constraints

Visualization

Related Problems

Common Approaches

Binary Search on Prefix Sums — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler