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

1. 描述

• Implement the RandomizedSet class:
  • RandomizedSet() Initializes the RandomizedSet object.
  • bool insert(int val) Inserts an item val into the set if not present. Returns true if the item was not present, false otherwise.
  • bool remove(int val) Removes an item val from the set if present. Returns true if the item was present, false otherwise.
  • int getRandom() Returns a random element from the current set of elements (it’s guaranteed that at least one element exists when this method is called). Each element must have the same probability of being returned.
  • You must implement the functions of the class such that each function works in average O(1) time complexity.
• 设计一个支持在平均 时间复杂度 O(1) 下，执行以下操作的数据结构。
  • insert(val)：当元素 val 不存在时，向集合中插入该项。
  • remove(val)：元素 val 存在时，从集合中移除该项。
  • getRandom：随机返回现有集合中的一项。每个元素应该有相同的概率被返回。

2. 分析

首先思考什么数据结构可以在O(1)时间插入删除——哈希表，但哈希表在getRandom时遇到问题，因为getRandom需要产生一个随机索引后根据此索引找到元素，但哈希表不具备索引（如创建indexToValue类型的哈希表，则无法常数时间删除）。
所以要想实现getRandom必须得使用能随机存取的数据结构，首先想到数组。但用数组存储的问题在于插入删除时需要移动元素，时间复杂度为O(n)，这里可以采用尾部插入、删除的方式解决，使插入、删除只需要O(1)时间。
对于插入操作，在数组尾部插入即可，但对于删除操作，需要在O(1)时间内找到待删除元素，这就需要用到一个valToIndex哈希表，在每次插入新元素时，存储元素值及其下标。执行删除时，先在哈希表中找到对应下标后，在数组中把元素交换到尾部进行删除，这样就实现了O(1)时间删除。注意删除元素后需要及时对哈希表进行更新。

3. 代码

class RandomizedSet {
    List<Integer> nums;
    Map<Integer, Integer> valToIndex;
    Random rand = new Random();
    /** Initialize your data structure here. */
    public RandomizedSet() {
    	// ArrayList底层为动态数组
        nums = new ArrayList<>();
        valToIndex = new HashMap<>();
    }
    
    /** Inserts a value to the set. Returns true if the set did not already contain the specified element. */
    public boolean insert(int val) {
        if(valToIndex.containsKey(val)) return false;
        valToIndex.put(val, nums.size());
        nums.add(nums.size(), val);
        return true;
    }
    
    /** Removes a value from the set. Returns true if the set contained the specified element. */
    public boolean remove(int val) {
        if(!valToIndex.containsKey(val)) return false;
        // 在哈希表中找到对应下标
        int idx = valToIndex.get(val);
        int lastElement = nums.get(nums.size() - 1);
        // 相当于与末尾元素交换后删除
        nums.set(idx, lastElement);
        // 对哈希表进行更新
        valToIndex.put(lastElement, idx);
        valToIndex.remove(val);
        nums.remove(nums.size() - 1);
        return true;
    }
    
    /** Get a random element from the set. */
    public int getRandom() {
        return nums.get(rand.nextInt(nums.size()));
    }
}

/**
 * Your RandomizedSet object will be instantiated and called as such:
 * RandomizedSet obj = new RandomizedSet();
 * boolean param_1 = obj.insert(val);
 * boolean param_2 = obj.remove(val);
 * int param_3 = obj.getRandom();
 */

4. 总结

数据结构的设计思路在于思考有哪些已知的数据结构能满足对特性的要求，必要时可对数据结构进行组合。