528. Random Pick with Weight

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the i t h i^{th} 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) = 0.25 (i.e., 25%), and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).
     
Example 1:
Example 2:

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 . l e n g t h < = 1 0 4 1 <= w.length <= 10^4 1<=w.length<=104
  • 1 < = w [ i ] < = 1 0 5 1 <= w[i] <= 10^5 1<=w[i]<=105
  • pickIndex will be called at most 1 0 4 10^4 104 times.

From: LeetCode
Link: 528. Random Pick with Weight


Solution:

Ideas:

1. Initialization (solutionCreate):

  • Compute a prefix sum array from the given weights. This allows us to efficiently map a random number to an index based on the weights.
  • Store this prefix sum and the total sum of weights for later use during the pick index operation.

2. Pick Index (solutionPickIndex):

  • Generate a random number within the range [1, total weight].
  • Use a binary search on the prefix sum array to find the smallest index where the prefix sum is greater than or equal to the random number. This step ensures that the probability of picking an index is proportional to its weight.

3 .Clean Up (solutionFree):

  • Free up the memory used by the solution object, including the dynamically allocated prefix sum array.
Code:
typedef struct {
    int* prefixSums;
    int totalWeight;
    int length;
} Solution;

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

int solutionPickIndex(Solution* obj) {
    int target = rand() % obj->totalWeight + 1;
    int low = 0, high = obj->length - 1;
    
    // Binary search to find the lowest index where the prefix sum is >= target
    while (low < high) {
        int mid = low + (high - low) / 2;
        if (obj->prefixSums[mid] < target)
            low = mid + 1;
        else
            high = mid;
    }
    
    return low;
}

void solutionFree(Solution* obj) {
    free(obj->prefixSums);
    free(obj);
}
04-18 15:30