Taro Logo

Insert Delete GetRandom O(1) - Duplicates allowed

Medium
5 views
2 months ago

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?

Sample Answer
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()

Naive Solution

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.

Optimal Solution

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.

Big(O) Run-time Analysis

  • Insert: O(1) average time complexity because appending to a list and adding to a dictionary takes constant time.
  • Remove: O(1) average time complexity because removing from a set and popping from a list takes constant time.
  • GetRandom: O(1) average time complexity because choosing a random element from a list takes constant time.

Big(O) Space Usage Analysis

  • The space complexity is O(n), where n is the number of elements in the collection, because we store the elements in a list and the indices in a dictionary.

Edge Cases

  • Removing an element that doesn't exist: The remove function returns False if the element is not in the collection.
  • Removing the last element: If the element to remove is the last element in the list, we can simply pop it from the list and remove it from the dictionary.
  • Removing an element with multiple occurrences: We only remove one occurrence of the element.
  • Empty Collection: The problem statement indicates that 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.