Generate Weighted Random Keys With Add And Remove

Question

The candidate is asked to design a system that generates weighted random keys, allowing for the addition and removal of keys with their associated weights.

Codemia · Accepted Answer

To solve the problem of generating weighted random keys with add and remove functionality, we can utilize a combination of a list and a cumulative weight array. The list will store the keys, while the cumulative weight array will help in selecting a key based on its weight efficiently.

1. **Data Structure**: Use a list to store the keys and a corresponding list to store cumulative weights. 
2. **Add Operation**: When adding a key, append it to the list and update the cumulative weights. 
3. **Remove Operation**: To remove a key, find its index, remove it from the list, and adjust the cumulative weights accordingly. 
4. **Random Selection**: Generate a random number between 0 and the total weight, and use binary search on the cumulative weights to find the corresponding key.

The time complexity for adding a key is O(1) for appending and O(n) for updating the cumulative weights. For removing a key, it is O(n) to find and remove it, and O(n) to update the weights. The random selection takes O(log n) due to binary search on the cumulative weights. Overall, the space complexity is O(n) for storing the keys and weights.

```python
import random
class WeightedRandom:
    def __init__(self):
        self.keys = []
        self.cumulative_weights = []
        self.total_weight = 0

def add(self, key, weight):
        self.keys.append(key)
        self.total_weight += weight
        self.cumulative_weights.append(self.total_weight)

def remove(self, key):
        if key in self.keys:
            index = self.keys.index(key)
            self.total_weight -= self.cumulative_weights[index]
            del self.keys[index]
            del self.cumulative_weights[index]
            for i in range(index, len(self.cumulative_weights)):
                self.cumulative_weights[i] -= self.cumulative_weights[index]

def get_random_key(self):
        rand_val = random.uniform(0, self.total_weight)
        low, high = 0, len(self.cumulative_weights) - 1
        while low < high:
            mid = (low + high) // 2
            if rand_val > self.cumulative_weights[mid]:
                low = mid + 1
            else:
                high = mid
        return self.keys[low]
```

Consider scenarios where weights are zero or negative, or when trying to remove a key that does not exist. These cases should be handled gracefully to avoid errors in the system.

Generate Weighted Random Keys With Add And Remove

Nuro

What the Interviewer Expects

Key Topics to Cover

How to Approach This

Possible Follow-up Questions

Sharpen Your Skills on Codemia

Sample Answer

Approach

Algorithm

Submit Your Answer

Nuro Software Engineer Interview Guide

Related Questions

Binary Search on graph

Count kth largest element in array

Find connected components in string

Merge Three Sorted strings

Binary Search on string