Design a driver pool management system that maintains a collection of available drivers and supports efficient random selection. The system must track drivers by unique IDs and names, allow adding and removing drivers, and provide uniform random selection of drivers from the active pool.
Your implementation should support three core operations:
Example 1:
Operations: pool = DriverPool() pool.insert(1, "Alice") → true pool.insert(2, "Bob") → true pool.getRandom() → "Alice" or "Bob" (50% each) pool.remove(1) → true pool.getRandom() → "Bob"
Example 2:
Operations: pool = DriverPool() pool.insert(10, "Charlie") → true pool.insert(10, "David") → false (ID 10 already exists) pool.remove(20) → false (ID 20 doesn't exist) pool.getRandom() → "Charlie" pool.remove(10) → true pool.getRandom() → null (pool is empty)
Example 3:
Operations: pool = DriverPool() pool.insert(5, "Eve") → true pool.insert(6, "Frank") → true pool.insert(7, "Grace") → true pool.remove(6) → true pool.getRandom() → "Eve" or "Grace" (50% each)
Hint 1: Data Structure Choice To achieve O(1) random access, you need to access elements by index. Consider using both a hash map for O(1) lookup by ID and a list for O(1) random access by index. How can you keep these two structures synchronized?
Hint 2: Efficient Removal Removing an element from the middle of a list is O(n). However, removing from the end is O(1). Can you swap the element to be removed with the last element before removing it? Don't forget to update your hash map accordingly.
Hint 3: What to Store in the Hash Map The hash map should map driver_id to something that helps you quickly locate and remove the driver. Consider storing the index position in the list along with the driver's name.
Full Solution `` Approach Explanation:
The solution uses two complementary data structures:
- Hash Map (id_to_data): Maps driver_id to a tuple of (name, index). This provides O(1) lookup to check if a driver exists and to find their position in the list.
- List (drivers): Stores tuples of (driver_id, name). This provides O(1) random access by index for the getRandom operation.
Key Operations:
Insert: Append to the end of the list (O(1)) and add to hash map (O(1)).
Remove: To avoid O(n) removal from the middle of the list, we use the swap-with-last technique:
- Find the driver's index from the hash map
- Swap it with the last element in the list
- Update the hash map for the swapped element
- Remove the last element (O(1))
- Delete from hash map
GetRandom: Generate a random index between 0 and len(drivers)-1, then return the name at that index.
Complexity Analysis:
- Time Complexity: O(1) average for all operations
- Space Complexity: O(n) where n is the number of drivers in the pool
This design ensures uniform probability for getRandom because each driver occupies exactly one position in the list, and we select from all positions with equal probability.
import random
class DriverPool:
def __init__(self):
# Hash map: driver_id -> (name, index in list)
self.id_to_data = {}
# List of (driver_id, name) tuples for O(1) random access
self.drivers = []
def insert(self, driver_id: int, name: str) -> bool:
"""
Insert a driver into the pool.
Returns False if driver_id already exists, True otherwise.
Time: O(1)
"""
if driver_id in self.id_to_data:
return False
# Add to the end of the list
index = len(self.drivers)
self.drivers.append((driver_id, name))
# Store driver_id mapping to (name, index)
self.id_to_data[driver_id] = (name, index)
return True
def remove(self, driver_id: int) -> bool:
"""
Remove a driver from the pool.
Returns False if driver_id doesn't exist, True otherwise.
Time: O(1)
"""
if driver_id not in self.id_to_data:
return False
# Get the index of the driver to remove
name, index = self.id_to_data[driver_id]
# Get the last driver in the list
last_index = len(self.drivers) - 1
last_driver_id, last_name = self.drivers[last_index]
# Swap the driver to remove with the last driver
self.drivers[index] = (last_driver_id, last_name)
# Update the hash map for the swapped driver
if last_driver_id != driver_id: # Only update if not removing the last element
self.id_to_data[last_driver_id] = (last_name, index)
# Remove the last element from the list
self.drivers.pop()
# Remove from hash map
del self.id_to_data[driver_id]
return True
def getRandom(self) -> str:
"""
Return a random driver's name with uniform probability.
Returns None if the pool is empty.
Time: O(1)
"""
if not self.drivers:
return None
# Pick a random index and return the name
random_index = random.randint(0, len(self.drivers) - 1)
driver_id, name = self.drivers[random_index]
return name