380. O(1) 时间插入、删除和获取随机元素

exiaohu 于 2022-04-13 发布

题目链接:380. O(1) 时间插入、删除和获取随机元素

很有意思的题目,哈希表实现常数时间复杂度的插入和删除,数组实现常数时间复杂度的随机选取。

因为只需要随机选值,因此数组中元素的顺序不重要,删除元素时,将该元素与队尾元素交换,然后删除队尾元素即可。使用哈希表维护每个元素的索引,以便于通过哈希表在数组中寻找到相应的位置。

import random


class RandomizedSet:
    def __init__(self):
        self.index = dict()
        self.data = list()

    def insert(self, val: int) -> bool:
        if val not in self.index:
            self.index[val] = len(self.data)
            self.data.append(val)
            return True
        return False

    def remove(self, val: int) -> bool:
        if self.data and val in self.index:
            index = self.index[val]
            if index == len(self.data) - 1:
                self.data.pop()
                del self.index[val]
                return True
            self.data[index], self.data[-1] = self.data[-1], self.data[index]
            self.data.pop()
            self.index[self.data[index]] = index
            del self.index[val]
            return True
        return False

    def getRandom(self) -> int:
        return random.choice(self.data)