Optimizing LRU Cache - Leetcode #146

Leetcode #146 asks us to design a data structure that contains the same functionalities of an Least Recently Used cache. In short, the LRU cache will store key value pairs up to a certain capacity, and will enforce that capacity by removing the least recently used key value pair.

Attempt 1: HashMap and LinkedList

In my first attempt, I chose to use a HashMap to store key value pairs (duh) and a LinkedList to simulate a queue.

As key/value pairs get used, the key gets added to the LinkedList. If the LinkedList already contains the key, then we remove it and readd it to the end. This way, the least recently used pair will always be at the head of the LinkedList, and we can pop it to enforce capacity.

Two helper methods are used to keep code clean:

1. updateCache(key): remove the key from the linkedlist and re-add it to the end. We can use the LinkedList.removeFirstOccurrence(key) method because

2. evict(): this is called when size of the hashmap equals that of the capacity. We pop the linkedlist to make room for a new pair to be added.

class LRUCache {

    int capacity; 
    LinkedList<Integer> usageHistory; 
    HashMap<Integer, Integer> pairs; 

    public LRUCache(int capacity) {
        this.capacity = capacity; 
        this.usageHistory = new LinkedList(); 
        this.pairs = new HashMap(); 
    }

    public int get(int key) {
        if(pairs.containsKey(key)){
            this.updateCache(key);
            return pairs.get(key); 
        }
        return -1; 
    }

    public void put(int key, int value) {
    //If contains key, capacity wont be breached (value will get updated)
        if((pairs.size()==capacity) && (!pairs.containsKey(key))){
            this.evict(); 
        }        
        pairs.put(key, value); 
        this.updateCache(key);
    }

    private void updateCache(int key){
        usageHistory.removeFirstOccurrence(key); 
        usageHistory.add(key); 
    }

    private void evict(){
        int keyToRemove = usageHistory.pop(); 
        pairs.remove(keyToRemove); 
    }
}

Although the above implementation passes a couple of the test cases, we receive a 'Time Limit Exceeded' error upon submission. This occurs when the test case has a capacity of 3000.

Inspecting my code, it looks clear to me that the slowest part of my code is in the updateCache(key) method, specifically when I call LinkedList's removeFirstOccurrence method. This method searches the linked list until it finds the key and runs in linear time. I guess O(n) isn't fast enough for this problem.

Attempt 2:

To speed up searching the linked list from linear time to constant time, I need to be able to have direct access to each linked list node. The idea behind this solution is to use a HashMap that contains the key as key, and the specific linkedlist node as value.

Unfortunately, there is no built-in class available in Java to represent these nodes, so we will need to implement it ourselves. My approach is to use a doubly linked list, where each node contains a key and a value. The key is necessary because without it, our LinkedList will have no way to access the HashMap.

To begin, I define custom classes for double linked list and for each node. For my get operations, I will need to remove a node from the linked list and add it to the tail. For my put operation, I will need to add a node to the tail, or if I need to evict, I need to find the head and remove it.

class DoubleLinkedList{

    Node fixedHead; 
    Node fixedTail; 

    public DoubleLinkedList(){
        fixedHead = new Node(0,0); 
        fixedTail = new Node(0,0); 
        fixedHead.next = fixedTail; 
        fixedTail.prev = fixedHead; 
    }

    //DoubleLinkedList Helper Methods 
    void moveToTail(Node node){
        //Remove Node first
        node.next.prev = node.prev;
        node.prev.next = node.next; 

        addToTail(node);        
    }
    void addToTail(Node node){
//Add to tail -> point fixedTail to new node
//Add new node's prev (oldTail) and next (fixedTail)
//Point old tail's next to new tail
        Node oldTail = fixedTail.prev; 
        fixedTail.prev = node; 
        node.prev = oldTail; 
        node.next = fixedTail; 
        oldTail.next = node; 
    }

    void removeHead(){
        Node newHead = fixedHead.next.next; 
        newHead.prev = fixedHead; 
        fixedHead.next = newHead; 
    }

    Node getHead(){
        return fixedHead.next; 
    }
}

class Node{
    int key;
    int val; 
    Node next; 
    Node prev; 
    Node(int key, int val){
        this.key = key; 
        this.val = val; 
    }
}

Replacing the built in LinkedList with a custom double linked list means we also need to adjust our original methods:

class LRUCache {

    int capacity; 
    DoubleLinkedList usageHistory; 
    HashMap<Integer, Node> pairs; 

    public LRUCache(int capacity) {
        this.capacity = capacity; 
        this.usageHistory = new DoubleLinkedList(); 
        this.pairs = new HashMap(); 
    }

    public int get(int key) {
        if(pairs.containsKey(key)){
            Node node = pairs.get(key); 
            this.updateCache(node);
            return node.val; 
        }
        return -1; 
    }

    public void put(int key, int value) {
        if(pairs.containsKey(key)){
            //Find node, update node value, call updateCahce()
            Node nodeToUpdate = pairs.get(key); 
            nodeToUpdate.val = value; 
            updateCache(nodeToUpdate); 

        } else {
            //Create new node
            Node newNode = new Node(key, value); 
            if(pairs.size()==capacity){
                this.evict(); 
            }
            //Create new node, add to map, then insert it at end of dll 
            pairs.put(newNode.key, newNode); 
            usageHistory.addToTail(newNode); 
        }
    }

    private void updateCache(Node node){
        //Remove Node from DoubleLinkedList and add to tail 
        usageHistory.moveToTail(node); 

    }

    private void evict(){
        //Node key to search and remove from HashMap
        Node nodeToRemove = usageHistory.getHead(); 
        int keyToRemove = nodeToRemove.key; 
        pairs.remove(keyToRemove); 
        usageHistory.removeHead(); 
    }
}

Now that I can search up and delete nodes using the HashMap, my put and get operations now run in O(1) time and are able to successfully pass all Leetcode test cases.

Bonus Method: LinkedHashMap

Java's LinkedHashMap class allows us to solve this problem without implementing our own doubly-linked list, as the class itself maintains a doubly-linked list that defines iteration ordering (default: order in which keys were inserted into the map). The only thing that we have to keep in mind is that re-insertion doesn't affect insertion order, so we need to remove and re-add the key for update operations.

class LRUCache {

    private LinkedHashMap<Integer, Integer> map;
    private int capacity;

    public LRUCache(int capacity) {
        map = new LinkedHashMap<>();
        this.capacity = capacity;
    }

    public int get(int key) {
        if(map.containsKey(key)) {
            int value = map.remove(key);
            map.put(key, value);
            return value;
        }
        return -1;
    }

    public void put(int key, int value) {
        if(map.containsKey(key)) {
            map.remove(key);
        }else if(map.size() == capacity) {
            map.remove(map.keySet().iterator().next());
        }
        map.put(key, value);
    }
}

As a final note, the LRU implementation in Android's source code also uses a LinkedHashMap implementation.