P188 支持随机获取元素的集合

描述#

这个题目说的是，你要设计一个增加版的集合（Set），它除了支持插入元素（insert）和删除元素（remove）的操作，还能等概率地随机获取当前集合中的元素（getRandom）。并且这 3 个操作的平均时间复杂度都要求是 O(1)。

以下是这个集合的使用示例：
// 初始化一个空集合
RandomizedSet set = new RandomizedSet();
// 成功插入数字 3，返回 true。set: [3]
set.insert(3);
// 3 已经存在于集合中，返回 false。set: [3]
set.insert(3);
// 成功插入数字 6，返回 true。set: [3, 6]
set.insert(6);
// 成功插入数字 9，返回 true。set: [3, 6, 9]
set.insert(9);
// 以 1/ 3 的概率返回 3，6 或 9。set: [3, 6, 9]
set.getRandom();
// 成功删除数字 3，返回 true。set: [6, 9]
set.remove(3);
// 8 不存在于集合中，返回 false。set: [6, 9]
set.remove(8);
// 以 1/2 的概率返回 6 或 9。set: [6, 9]
set.getRandom();

代码#

public class AlgoCasts {

  public class RandomizedSet {

    private final Map<Integer, Integer> map = new HashMap<>();
    private final List<Integer> data = new ArrayList<>();
    private final Random RANDOM = new Random();

    public boolean insert(int val) {
      if (map.containsKey(val)) return false;
      data.add(val);
      map.put(val, data.size()-1);
      return true;
    }

    public boolean remove(int val) {
      if (!map.containsKey(val)) return false;
      int idx = map.get(val);
      int lastVal = data.get(data.size()-1);
      data.set(idx, lastVal);
      data.remove(data.size()-1);
      map.put(lastVal, idx);
      map.remove(val);
      return true;
    }

    public int getRandom() {
      int idx = RANDOM.nextInt(data.size());
      return data.get(idx);
    }
  }

}