LeetCode 381 O（1）时间内插入，删除和获取元素

设计一个支持在平均 时间复杂度 O(1) 下， 执行以下操作的数据结构。

注意: 允许出现重复元素。

insert(val)：向集合中插入元素 val。
remove(val)：当 val 存在时，从集合中移除一个 val。
getRandom：从现有集合中随机获取一个元素。每个元素被返回的概率应该与其在集合中的数量呈线性相关。

from collections import defaultdict
import random


class RandomizedCollection:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.val2idx = defaultdict(lambda: set())
        self.nums = []

    def insert(self, val: int) -> bool:
        """
        Inserts a value to the collection. Returns true if the collection did not already contain the specified element.
        """
        self.nums.append(val)
        size = len(self.nums)
        self.val2idx[val].add(size - 1)
        return len(self.val2idx[val])==1

    def remove(self, val: int) -> bool:
        """
        Removes a value from the collection. Returns true if the collection contained the specified element.
        """
        if len(self.val2idx[val]) == 0:
            return False
        idx1 = self.val2idx[val].pop()
        size = len(self.nums)
        if size != idx1 + 1:
            val1 = self.nums[size - 1]
            self.nums[idx1], self.nums[size - 1] = self.nums[size - 1], self.nums[idx1]
            self.nums.pop()
            self.val2idx[val1].remove(size - 1)
            self.val2idx[val1].add(idx1)
        else:
            self.nums.pop()
        return True

    def getRandom(self) -> int:
        """
        Get a random element from the collection.
        """
        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()