algoadvance

Leetcode 528. Random Pick with Weight

Problem Statement

Given an array w of positive integers where w[i] describes the weight of index i, write a function pickIndex that randomly picks an index in proportion to its weight.

Implement the Solution class:

class Solution {
public:
    Solution(vector<int>& w);
    int pickIndex();
};

Clarifying Questions

  1. Input Constraints:
    • What is the range of values that w can take?
    • What is the range of the length of w?
  2. Randomness:
    • Is there a specific random seed we should use, or do we assume truly random behavior?
  3. Performance:
    • Are there any performance constraints or requirements, especially regarding time complexity for initialization and pick operations?

Code

#include <vector>
#include <cstdlib>

class Solution {
public:
    Solution(std::vector<int>& w) : prefixSums(w.size(), 0) {
        // Calculate the prefix sums
        prefixSums[0] = w[0];
        for (int i = 1; i < w.size(); ++i) {
            prefixSums[i] = prefixSums[i - 1] + w[i];
        }
    }
    
    int pickIndex() {
        int totalSum = prefixSums.back();
        int randWeight = rand() % totalSum;

        // Binary search for the correct index
        int low = 0, high = prefixSums.size() - 1;
        while (low < high) {
            int mid = low + (high - low) / 2;
            if (randWeight >= prefixSums[mid]) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }
        return low;
    }

private:
    std::vector<int> prefixSums;
};

Strategy

  1. Initialization (Solution constructor):
    • Calculate prefix sums of the input weights. This enables us to map the weights to ranges, which will facilitate the index selection later.
    • prefixSums[i] will hold the sum of w[0] to w[i].
  2. Index Picking (pickIndex method):
    • Generate a random number between 0 and the total sum of weights minus one.
    • Use binary search to find the smallest index such that the prefix sum is greater than the random number. This index corresponds to the weight ranges correctly.
    • The binary search ensures that the “search” part of picking an index is efficient.

Time Complexity

This solution ensures that the pickIndex operation is efficient even for large input sizes, striking a balance between preprocessing time and query time.

Cut your prep time in half and DOMINATE your interview with AlgoAdvance AI