Random Pick with Weight

 

1. Clarify the problem: Before diving into solving the problem, let's clarify the requirements:

  • We are given an array of positive integers representing the weights of elements.
  • We need to design a data structure that allows us to efficiently pick an element randomly with probability proportional to its weight.

2. Analyze the problem: To solve this problem, we can precalculate the cumulative sum of the weights and use binary search to find the index corresponding to the randomly picked weight.

3. Design an algorithm: Here is the algorithm to solve the problem:

  1. Initialize a cumulative sum array, prefixSum, to store the cumulative sums of the weights.
  2. Initialize a variable totalSum to store the sum of all weights.
  3. Iterate through the given weights and calculate the cumulative sum, updating prefixSum and totalSum.
  4. Generate a random number between 1 and totalSum.
  5. Use binary search to find the index i corresponding to the generated random number in the prefixSum array.
    • If the random number is less than or equal to prefixSum[i], return the index i.
    • If the random number is greater than prefixSum[i], perform binary search on the right half of the prefixSum array.
  6. Return -1 if the cumulative sum array is empty.

4. Explain your approach: The approach involves precalculating the cumulative sum of the weights and generating a random number between 1 and the total sum of weights. Using binary search, we find the index corresponding to the generated random number in the cumulative sum array. The index represents the randomly picked element with probability proportional to its weight.

5. Write clean and readable code:

python
import random class Solution: def __init__(self, w): self.prefixSum = [] self.totalSum = 0 # Calculate the cumulative sum for weight in w: self.totalSum += weight self.prefixSum.append(self.totalSum) def pickIndex(self): # Generate a random number between 1 and totalSum randomNum = random.randint(1, self.totalSum) # Binary search to find the index left, right = 0, len(self.prefixSum) - 1 while left < right: mid = left + (right - left) // 2 if randomNum <= self.prefixSum[mid]: right = mid else: left = mid + 1 return left # Test case weights = [1, 3, 2, 4] solution = Solution(weights) index = solution.pickIndex() print(index)

6. Test your code: Let's test the code with the given test case:

python
weights = [1, 3, 2, 4] solution = Solution(weights) index = solution.pickIndex() print(index)

The expected output is a random index from 0 to 3, each with probability proportional to its weight.

7. Optimize if necessary: The current solution is already efficient with a time complexity of O(log n) for each pickIndex operation due to binary search. There are no further optimizations needed for this problem.

8. Handle error cases: The code assumes that the input array w contains positive integers. If the array is empty or contains non-positive integers, the behavior of the code may not be well-defined.

9. Discuss complexity analysis:

  • The time complexity of constructing the data structure is O(n), where n is the length of the input array w.
  • The time complexity of the pickIndex operation is O(log n) due to binary search.
  • The space complexity is O(n) to store the cumulative sum array prefixSum.
Next Post Previous Post