🎲

[leetCode] Top Interview 150: Insert Delete GetRandom O(1)

2024/01/09に公開

Java

LeetCode

tech

リンク

概要

RandomizedSetクラスを実装していく。
- bool insert(int val): データを挿入する。成功ならtrue, ダメならfalseを返す。
- bool remove(int val): データを削除する。成功ならtrue, ダメならfalseを返す。
- int getRandom(): ランダムにひとつデータを取得して返す。
各関数の平均計算量が $O(1)$ になるように実装する。

解法

ぱっと思いつくのは、 HashSet クラスを使用した実装。

class RandomizedSet {
    private Set<Integer> set;

    public RandomizedSet() {
        set = new HashSet<>();    
    }
    
    public boolean insert(int val) {
        return this.set.add(val);
    }
    
    public boolean remove(int val) {
        return this.set.remove(val);
    }
    
    public int getRandom() {
    int size = this.set.size();
    int index = (int) (Math.random() * size);
    Iterator<Integer> iterator = this.set.iterator();
    while (index-- > 0) {
        iterator.next();
    }
        return iterator.next();
    }
}

これでも一応通る...けど、getRandom() の処理内でwhile を使用しているため、平均計算量が $\text{$\frac N 2$}$ になってしまう。
前提として、 HashSet にはインデックス参照による要素の取得がない。 そのため、乱数で適当にインデックスを決めて取得、みたいなことが難しい。

この問題を解決するためには、ArrayListで要素を管理しつつ、HashMapで要素の位置をメモし、検索処理を軽量化する必要がある。ついでに、HashMap.containsKey()で要素の重複をチェックすれば、かんたんにHashSetの性質を引き継ぐことができる。
これを踏まえて実装したのが以下のコード。

import java.util.*;

class RandomizedSet {
    private Map<Integer, Integer> map;
    private List<Integer> list;
    private Random random;

    public RandomizedSet() {
        map = new HashMap<>();
        list = new ArrayList<>();
        random = new Random();
    }

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

    public boolean remove(int val) {
        if (!map.containsKey(val)) {
            return false;
        }
        int indexToRemove = map.get(val);
        int lastElement = list.get(list.size() - 1);

        list.set(indexToRemove, lastElement);
        map.put(lastElement, indexToRemove);

        list.removeLast();
        map.remove(val);
        return true;
    }

    public int getRandom() {
        return list.get(random.nextInt(list.size()));
    }
}

注意点として、ArrayList.remove()の計算量は $O(1)$ ではない。 要素を削除したあと、その要素の右にあるすべての要素を左へと１つずらすためである。
これは、 Java の公式ドキュメントからも記述が確認できる。

例外として、末尾要素の削除は $O(1)$ であることが保証される。そのため、いちど削除する要素を末尾にずらしてから削除している。

おわりに

普段脳死でコレクションを使っていたので「HashSetでアレコレすればいけるやろｗ」と思っていたらハマりました。学び。

リンク

概要

解法

おわりに

Discussion