C++でO(1)で読み書きできるLRUキャッシュを実装する

2020-04-26 c++

最初次のようなコードを書いたがタイムアウトしてしまった。 C++のstd::mapは二分木で実装されているので各操作にO(log n)かかり、dequeの先頭への挿入はO(1)でできるが重複するキーを探して削除するのにO(n^2)かかってしまう。

C++ STLのContainersとAlgorithms - sambaiz-net

class LRUCache {
public:
    deque<int> Q; // key
    map<int, int> V; // <key, value>
    int cap = 0;
    LRUCache(int capacity) {
        cap = capacity;
    }
    
    int get(int key) {
        if (V.count(key) != 0) {
            for (int i = 0; i < Q.size(); i++) {
                if (Q[i] == key) {
                    Q.erase(Q.begin() + i);
                    break;
                }
            }
            Q.push_front(key);
            return V[key];
        }
        return -1;
    }
    
    void put(int key, int value) {
        if (cap == 0) {
            return;
        }
        if (get(key) == -1) {
            if (Q.size() == cap) {
                V.erase(Q.back());
                Q.pop_back();
            }
            Q.push_front(key);
        }
        V[key] = value;
    }
};

ヒントにO(1)で実装できると書いてあったので、mapの代わりにO(1)で読み書きできるunordered_map (hash map)を、dequeの代わりに削除もO(1)でできるlistを使うことにした。あとは普通にlistにアクセスするとO(n)になるのをどうするかだが、iteratorを持っておけばO(1)で直接アクセスできる。これで読み書き共にO(1)となり時間内に終わるようになった。

class LRUCache {
public:
    list<tuple<int, int>> Q; // {key, value}
    unordered_map<int, list<tuple<int, int>>::iterator> iter; // <key, iterator>
    int cap = 0;
    LRUCache(int capacity) {
        cap = capacity;
    }
    
    int get(int key) {
        if (iter.count(key) == 0) {
            return -1;
        }
        auto it = iter[key];
        auto [k, v] = *it;
        Q.erase(it);
        Q.push_front({key, v});
        iter[key] = Q.begin();
        return v;
    }
    
    void put(int key, int value) {
        if (iter.count(key) != 0) {
            Q.erase(iter[key]);
        }
        Q.push_front({key, value});
        iter[key] = Q.begin();
        if (Q.size() > cap) {
            auto [k, v] = *(--Q.end());
            iter.erase(k);
            Q.pop_back();
        }
    }
};