leetcode(380). Insert Delete GetRandom O(1)

problem

Design a data structure that supports all following operations in
average O(1) time.

1. insert(val): Inserts an item val to the set if not already present.
2. remove(val): Removes an item val from the set if present.
3. getRandom: Returns a random element from current set of elements. Each element must have the same probability of being returned.

分析

使用普通的hash集合就可以满足前两项要求，但是对于随机取出一个元素来说可能需要使用一个支持索引的数据结构，这样就可以生成一个随机整数来随机取出一个元素。

但是使用线性表的问题就是在删除时需要移动后面的元素，对此我们使用一个val-position的dict维持元素和位置的映射，在删除时只把最后一个元素移到被删除的位置。

import random

class RandomizedSet(object):

    def __init__(self):
        self.nums, self.pos = [], {}

    def insert(self, val):
        if val not in self.pos:
            self.nums.append(val)
            self.pos[val] = len(self.nums) - 1
            return True
        return False


    def remove(self, val):
        if val in self.pos:
            idx, last = self.pos[val], self.nums[-1]
            self.nums[idx], self.pos[last] = last, idx
            self.nums.pop(); self.pos.pop(val, 0)
            return True
        return False

    def getRandom(self):
        return random.choice(self.nums)

总结

dict实现了val到pos的映射，可以用来实现查找和删除，list实现了pos到val的映射，可以用来返回随机元素。

这个数据结构实现的关键就是在删除元素时维持这种关系，通过被删除元素的值找到它在list中的位置，然后使用最后一个元素覆盖掉这个位置，同时修改dict中的数据。这样的操作时间复杂度为$O(1)$

ps：再思考这种问题时应该思路清晰一点，比如考虑删除操作时只考虑删除一个元素的操作，不用想得太多，只要单个元素能完成，更多的也一样。