LFU by tree map [solution]

  • Requirement
    • Remove less frequently used items by key;
    • If the most frequently used items has more than one, remove the eldest one;
  • Consideration
    • when get/put operated, we should touch the item
      • update timestamp;
      • update frequency;
  • Implementation remarks
    • Use record type for item;
    • Use TreeMap for ordered handling:
      • pollFirstEntry()
      • Need to re-put entry when frequency and timestamp changed
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;

record Node(int value, int stamp, int freq) {}

public class LFUCache {
    private int seq = 0, capacity;
    private Map<Integer, Node> dict;
    private TreeMap<Node, Integer> orderedMap;

    public LFUCache(int capacity) {
        this.capacity = capacity;
        dict = new HashMap<>();
        orderedMap = new TreeMap<>((a, b) -> a.freq() == b.freq() ? a.stamp() - b.stamp() : a.freq() - b.freq());
    }

    public int get(int key) {
        if (!dict.containsKey(key)) return -1;
        Node old = dict.get(key);
        orderedMap.remove(old);
        Node neo = new Node(old.value(), ++seq, old.freq() + 1);
        dict.put(key, neo);
        orderedMap.put(neo, key);
        return neo.value();
    }

    public void put(int key, int val) {
        if (dict.containsKey(key)) {
            Node old = dict.get(key);
            orderedMap.remove(old);
            dict.remove(key);
            Node neo = new Node(val, ++seq, old.freq() + 1);
            dict.put(key, neo);
            orderedMap.put(neo, key);
        } else {
            if (dict.size() == capacity)
                dict.remove(orderedMap.pollFirstEntry().getValue());
            Node neo = new Node(val, ++seq, 1);
            dict.put(key, neo);
            orderedMap.put(neo, key);
        }
    }
}