使用 Java 设计一种数据结构实现O(1) 时间插入、删除和获取随机元素

一、前言

今天从 LeetCode 上看到的看到一道题目挺有意思的，写个题解记录一下。

设计一个数据结构，要求在O(1) 时间插入、删除和获取随机元素；
实现RandomizedSet 类：
RandomizedSet() 初始化 RandomizedSet 对象
bool insert(int val) 当元素 val 不存在时，向集合中插入该项，并返回 true ；否则，返回 false 。
bool remove(int val) 当元素 val 存在时，从集合中移除该项，并返回 true ；否则，返回 false 。
int getRandom() 随机返回现有集合中的一项（测试用例保证调用此方法时集合中至少存在一个元素）。每个元素应该有 相同的概率 被返回。
你必须实现类的所有函数，并满足每个函数的 平均 时间复杂度为 O(1) 。
原题链接：https://leetcode.cn/problems/insert-delete-getrandom-o1/description/?envType=study-plan-v2&envId=top-interview-150

二、题目分析

让我们提炼一下上面这个问题中的关键点：

1. 在插入、删除时都要先判断元素是否存在，再进行插入、删除操作，数组（O(n)）、链表（O(n)）结构时间明显都不能满足要求，我们很容易想到使用哈希表（O(1)）;
2. 随机获取元素，且要求每个元素应该有相同的概率被返回，这个哈希表就不能满足了。随机获取集合中的一个元素，我们容易想到数组，数组是连续的内存天然支持随机访问，我们每次可以获取一个集合长度内的随机数下标 index，根据 index 我们就可以在 O(1) 时间内返回元素了；

单哈希表或数组都不能满足上述题目所有要求，我们再看下能否用空间换时间，同时使用哈希表+数组来实现 O(1) 时间插入、删除和获取随机元素呢？
答案是肯定的，哈希表本身就满足上面分析关键点中的1，我们可以同时存一份集合数据在数组，在关键点2中要求随机访问时我们根据随机获取的下标从数据中获取元素就可以了。

我可真是个小天才

啊不对，是不是忘记了什么，数组中备份一份数据，随机访问倒是妥妥的，插入每次都在最后用动态数据也没问题了，可删除动作时备份数组中也需要删除，如果从中间节点删除会涉及到后续所有数据移动，啊这我的 O(1) 不是凉的透透的。

诶，有个生活小妙招，我们本身就使用了哈希表，如果我们在哈希表中同时维护每个元素在备份数组中的下标 index, 抽象来看删除对备份数组来说本质就只是去掉该元素，然后长度减1即可。那我们备份数组中删除不是可以直接把最后一个元素移动到要删除元素位置上，这样的好处是不会造成删除元素后续所有元素的移动。嗯非常完美！

结论：动态数组（ArrayList）+ 哈希表（HashMap）的组合数据结构，可以解决我们 O(1) 时间插入、删除和获取随机元素的需求。

三、上代码

import java.util.*;

/**
 * O(1) 时间插入、删除和获取随机元素
 */
public class RandomizedSet {

    /**
     * 哈希表，key为元素，value为元素在动态数据中的下标index
     */
    private Map<Integer, Integer> hashMap;

    /**
     * 动态数组
     */
    private List<Integer> arrayList;

    private Random random;

    /**
     * 初始化 RandomizedSet 对象
     */
    public RandomizedSet() {
        hashMap = new HashMap<>();
        arrayList = new ArrayList<>();
        random = new Random();
    }

    /**
     * 当元素 val 不存在时，向集合中插入该项，并返回 true ；否则，返回 false
     * @param val
     * @return
     */
    public boolean insert(int val) {
        if (hashMap.containsKey(val)) {
            //元素已存在
            return false;
        }
        
        //新插入元素在动态数组下标，即 arrayList.size() - 1 + 1
        int index = arrayList.size();
        hashMap.put(val, index);
        arrayList.add(val);
        return true;
    }

    /**
     * 当元素 val 存在时，从集合中移除该项，并返回 true ；否则，返回 false
     * @param val
     * @return
     */
    public boolean remove(int val) {
        if (!hashMap.containsKey(val)) {
            //元素已存在
            return false;
        }

        //待删除元素下标
        int index = hashMap.get(val);
        //当前动态数据最后一个元素
        int lastElement = arrayList.get(arrayList.size() - 1);
        //把当前最后一个元素移动到删除元素下标位置
        arrayList.set(index, lastElement);
        //更新移动元素在哈希表中的下标映射
        hashMap.put(lastElement, index);

        //在哈希表、动态数组中删除要删的元素
        hashMap.remove(val);
        //注意，数组中原本的最后一个元素已移动到删除元素位置，因此直接将动态数组现在的尾部内存空间释放即可
        arrayList.remove(arrayList.size() - 1);
        return true;
    }

    /**
     * 随机返回现有集合中的一项（测试用例保证调用此方法时集合中至少存在一个元素）。每个元素应该有 相同的概率 被返回。
     * @return
     */
    public int getRandom() {
        int randomIndex = random.nextInt(arrayList.size());
        return arrayList.get(randomIndex);
    }
    
}

四、结尾

以上就是这个问题的一种解法，大家要是有其他好的方案也欢迎讨论。