Generate Weighted Random Keys With Add And Remove
Company: Nuro
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
Design a class that stores string keys with positive integer weights. `add_key(key, weight)` inserts or updates a key. `generate()` returns one key randomly with probability proportional to its current weight. In the follow-up, add `remove_key(key)` and discuss how to support updates and removals efficiently.
Function signature:
```python
class WeightedKeyGenerator:
def add_key(self, key: str, weight: int) -> None:
pass
def remove_key(self, key: str) -> None:
pass
def generate(self) -> str:
pass
```
Constraints:
- Weights are positive integers.
- Calling `generate` on an empty generator is invalid.
- The base solution may use prefix sums and binary search.
- The follow-up should discuss a segment tree or Fenwick tree for efficient updates.
- Randomness should be proportional to weight.
Examples:
```text
keys: a weight 1, b weight 3. generate should return b about three times as often as a over many trials.
```
Quick Answer: Review a weighted random key-generation coding prompt that introduces add, generate, and remove operations. The interview topic is useful for discussing prefix sums, binary search, segment trees or Fenwick trees, update costs, and deterministic testing limits for randomized behavior.
Implement `simulate_weighted_keys(operations)` for string keys with positive integer weights. In a real generator, `generate()` samples a uniform integer ticket from `1` through the current total weight. The console receives that ticket explicitly so execution is deterministic while testing the same proportional-selection logic.
Process these operations:
- `["add", key, weight]` inserts a key or replaces its current weight.
- `["remove", key]` removes a key; an absent key has no effect.
- `["generate", ticket]` returns the key whose cumulative weight interval contains `ticket`.
Keys are ordered by their first-ever add operation. Updating, removing, or re-adding a key preserves that position. Return one key for each generate operation, in order. Use a Fenwick tree or segment tree so updates, removals, and ticket lookup are logarithmic.
Constraints
- 0 <= len(operations) <= 200,000.
- Keys are non-empty strings.
- Every add weight is a positive integer.
- add replaces the current weight when the key is already active.
- remove on an absent key has no effect.
- Every generate operation occurs while at least one key is active.
- For generate, 1 <= ticket <= the sum of all current weights.
- A uniformly random ticket produces selection probability weight / total_weight.
- A removed and later re-added key keeps its first-ever ordering position.
Examples
Input: ([],)
Expected Output: []
Input: ([["add", "a", 1], ["add", "b", 3], ["generate", 1], ["generate", 2], ["generate", 4]],)
Expected Output: ['a', 'b', 'b']
Hints
- Treat each positive weight as a contiguous interval in the global ticket range.
- A point-update prefix-sum structure can apply weight deltas without rebuilding every prefix.
- Ticket lookup asks for the smallest prefix sum that is at least the ticket.
- Removal is a weight update to zero; adding an existing key applies the difference from its old weight.