LeetCode //C - 528. Random Pick with Weight

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^{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^4$
• $1<=w[i]<=10^5$
• pickIndex will be called at most $10^4$ 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);
}