Implement a data structure RandomizedCollection
that supports inserting, removing, and getting a random element in O(1) average time complexity, even with duplicates.
RandomizedCollection()
Initializes the empty RandomizedCollection
object.bool insert(int val)
Inserts an item val
into the multiset, even if the item is already present. Returns true
if the item is not present, false
otherwise.bool remove(int val)
Removes an item val
from the multiset if present. Returns true
if the item is present, false
otherwise. Note that if val
has multiple occurrences in the multiset, we only remove one of them.int getRandom()
Returns a random element from the current multiset of elements. The probability of each element being returned is linearly related to the number of the same values the multiset contains.For example:
RandomizedCollection randomizedCollection = new RandomizedCollection();
randomizedCollection.insert(1); // return true
randomizedCollection.insert(1); // return false
randomizedCollection.insert(2); // return true
randomizedCollection.getRandom(); // Returns 1 with probability 2/3 or 2 with probability 1/3.
randomizedCollection.remove(1); // return true
randomizedCollection.getRandom(); // Returns 1 or 2, both equally likely.
How would you implement this data structure?
import random
class RandomizedCollection:
def __init__(self):
self.nums = [] # Store the numbers
self.indices = {} # Store indices of each number
def insert(self, val: int) -> bool:
self.nums.append(val)
if val in self.indices:
self.indices[val].add(len(self.nums) - 1)
return False
else:
self.indices[val] = {len(self.nums) - 1}
return True
def remove(self, val: int) -> bool:
if val not in self.indices:
return False
index_to_remove = self.indices[val].pop()
if not self.indices[val]:
del self.indices[val]
if index_to_remove == len(self.nums) - 1:
self.nums.pop()
return True
last_element = self.nums.pop()
self.nums[index_to_remove] = last_element
self.indices[last_element].add(index_to_remove)
self.indices[last_element].remove(len(self.nums))
return True
def getRandom(self) -> int:
return random.choice(self.nums)
# Your RandomizedCollection object will be instantiated and called as such:
# obj = RandomizedCollection()
# param_1 = obj.insert(val)
# param_2 = obj.remove(val)
# param_3 = obj.getRandom()
A naive solution would involve using a list to store the elements and iterating through it for removal. However, this would result in O(n) time complexity for removal, which doesn't meet the requirement of O(1) average time complexity.
The optimal solution uses a list to store the elements and a dictionary to store the indices of each element in the list. This allows for O(1) average time complexity for insertion, removal, and getting a random element.
remove
function returns False
if the element is not in the collection.getRandom
will only be called when there is at least one item in the collection. Thus, we don't need to explicitly handle this edge case.