Implement an LRU Cache

Problem: Implement an LRU cache. An LRU cache stores items up to its capacity. If the capacity is exceeded, it removes the least recently used item.

Constraints:

  • Storing and accessing items in the cache must be done in O(1) time.
  • The capacity of the cache is between 1 and 5000.
Table of contents

Initial Thought Process

In general, it is important to clarify the requirements and constraints before jumping into implementation. That is especially true for data structure design questions such as this, because such questions want to test your ability to write logical and maintainable code. The interviewer is likely to leave out many details, and it is your responsibility to clarify what you need to do.

Requirements and Constraints

Data type

First of all, let us clarify what kind of items the LRU cache will store. There can be many choices, but let us assume that the we agree with the interviewer that the cache will store 32-bit signed integer for the keys and values.

API

It is a good idea to write down the list of APIs and their specification details before implementing them. Doing so will not only guide your through the implementation, but also make it easier for the interviewer to understand you. Since we are dealing with a cache, it is natural to have get and put APIs. This idea leads to the following APIs:

  1. LRUCache(int capacity): This will be the initializer for the data structure. It creates an LRU cache with the given capacity.
  2. get(Integer key): It will retrieve the value for the given key from the cache.
  3. put(Integer key, Integer val): It will store the given key value pairs into the cache.

A put API call will remove the least recently used item from the cache, if the capacity is exceeded. It is a good idea to clarify what “used” means here. Let us say that an item is used whenever there is a get or put API call for its key.

There is one edge case to which we need to pay attention. What is the behavior of a get API call, if the given key does not exist in the cache? For this question, let us say that we will return null. As you can see, listing out the APIs helps us think about the possible edge cases.

Thread-safety

There can be a race condition if two threads simultaneously write an item into the cache. Depending on how we store the data, such concurrent accesses can result in a corrupted state. Does our data structure need to be thread-safe? It is important to clarify the behavior upfront. For simplicity, let us suppose that the requirement does not require the cache to be thread-safe. Ideas for safe concurrency could make a good follow-up question.

Intuition

Let us start with the following observations:

  1. We need some sort of map to be able to access an item by its key in a constant time.
  2. The cache needs to maintain an order of the items in terms of latest time that they are used.

The implication of the first observation is rather clear: the cache can internally put items in a map data structure. What about the second observation of having to maintain an order of data access time? One way to do it would be to maintain a doubly-linked list.

Let us say we have the following linked list of items where the head is the most recently used item, and the tail is the least recently used item.

[A] <-> [B] <-> [C] <-> [D] <-> [E]
 ^                               ^
(Most recently used)         (Least recently used)

If item D is accessed, we can update the order of items by bringing it to the head. Because all nodes are doubly-linked, we can remove a node from the list, and insert it to the head, all in a constant time.

[D] <-> [A] <-> [B] <-> [C] <-> [E]
 ^                               ^
(Most recently used)         (Least recently used)

If a new item, F, is inserted to the cache, and the cache’s capacity is exceeded, we can evict the least recently used item, E, by updating the tail. Specifically, the new tail will be the previous node of E, namely, C.

[F] <-> [D] <-> [A] <-> [B] <-> [C]
 ^                               ^
(Most recently used)         (Least recently used)

Implementation

The observations lead us to the following starting point. The cache will have a hash map and a doubly-linked list. Also, it will remember the head and tail, so that it can efficiently update the order in the linked list. The hash map will point to a node in the linked list. The values will be stored in the nodes.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

class LRUCache {
    private int capacity;

    private class Node {
        public Node next;
        public Node prev;
        public Integer val;
        public Integer key;

        public Node(Integer key, Integer val) {
            this.key = key;
            this.val = val;
        }
    }
    private Node head;
    private Node tail;

    private Map<Integer, Node> data;

    public LRUCache(int capacity) {
        if (capacity < 1 || capacity > 5000) {
            throw new IllegalArgumentException(
                "A capacity must be between 1 and 5000."
            );
        }
        this.capacity = capacity;
        this.data = new HashMap<>();
    }

    public Integer get(Integer key) {
    }

    public void put(Integer key, Integer val) {
    }
}

The get and put implementation would need to bring a node to the head. So, let us implement moveToFront to reduce duplication, and break down the logic into maintainable pieces. Moving a node to the front is a two-step process of (1) removing the node from the linked list, and (2) inserting the node at the head. Note that the first process of removing the node from the list can also be re-used in the eviction logic.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

class LRUCache {
    private void moveToFront(Node node) {
        // Node is already head
        if (node.prev == null) {
            return;
        }

        removeNode(node);
        insertHead(node);
    }

    private void removeNode(Node node) {
        Node prev = node.prev;
        Node next = node.next;

        // Pull node out of the current position.
        prev.next = next;
        if (next != null) {
            next.prev = prev;
        }
    }

    private void insertHead(Node node) {
        node.prev = null;

        // If the node is the first item in the cache.
        if (head == null) {
            head = node;
            tail = node;
        } else {
            head.prev = node;
            node.next = head;
            head = node;
        }
    }
}

There is one important edge case we will need to handle while removing a node. That edge case is when the node being removed is the tail. Since the cache needs to remember the current tail, it will need to update the reference to the tail. Since we have a lot of moving parts and edge cases to consider, let us define a separate method updateTail to help us understand and maintain the logic.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

class LRUCache {
    private void removeNode(Node node) {
        Node prev = node.prev;
        Node next = node.next;

        // Pull node out of the current position.
        prev.next = next;
        if (next != null) {
            next.prev = prev;
        }

        // If the node is a tail, we need to update the tail
        if (node == tail) {
            updateTail();
        }
    }

    private void updateTail() {
        Node newTail = tail.prev;
        if (newTail != null) {
            newTail.next = null;
        }
        tail.prev = null;
        tail = newTail;
    }
}

Note that we do not need to handle a case in which a node being removed is the head. Since the cache will always have a capacity of at least 1, the head will never be evicted. Also, our moveToFront is a no-op if the node is the head.

Based on these methods, get API can be implemented in the following way:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class LRUCache {
    public Integer get(Integer key) {
        if (!data.containsKey(key)) {
            return null;
        }

        Node node = data.get(key);
        moveToFront(node);

        return node.val;
    }
}

And put can be implemented in the following way:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class LRUCache {
    public void put(Integer key, Integer val) {
        if (data.containsKey(key)) {
            // If the cache already contains the key,
            // (1) update the value in the node.
            // (2) insert the node at the front.
            Node node = data.get(key);
            node.val = val;

            moveToFront(node);
        } else {
            // If the cache doesn't have the key, insert the node at the front.
            Node node = new Node(key, val);
            data.put(key, node);

            insertHead(node);

            // Evict the least recently used item if the capacity is exceeded.
            if (data.size() > capacity) {
                evict();
            }
        }
    }
}

evict logic can be implemented by re-using the removeNode building block that we defined.

1
2
3
4
5
6
7
8

class LRUCache {
    private void evict() {
        Integer key = tail.key;

        removeNode(tail);
        data.remove(key);
    }
}

Complexity

Let us say N is the number of items in the cache.

  • The time complexity of get and put is expected to be O(1) assuming a perfect hash function, because we use a hash map to store and retrieve items.
  • The space complexity of the data structure is O(N) because the size of the linked list and the hash map will grow proportionally to the number of items.

Simplification (Sentinel Nodes)

In the previous implementation we constantly needed to check for edge cases such as removing tail, or when the node is the first item in the cache. If we can make this edge cases irrelevant, our implementation will be simpler and become easier to maintain and understand. The edge cases arose from the fact that we needed to change the head or tail. What if we do not have to change them at all?

We can maintain placeholder head and tail nodes ("sentinel nodes") that will never change. Then, instead of evicting the tail, we will evict the previous node of the sentinel tail. And, instead of inserting a node to the head, we will insert the node next to the sentinel head.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

class LRUCache {
    private int capacity;

    private class Node {
        public Node next;
        public Node prev;
        public Integer val;
        public Integer key;

        public Node(Integer key, Integer val) {
            this.key = key;
            this.val = val;
        }
    }
    private Node head;
    private Node tail;

    private Map<Integer, Node> data;

    public LRUCache(int capacity) {
        if (capacity < 1 || capacity > 5000) {
            throw new IllegalArgumentException(
                "A capacity must be between 1 and 5000."
            );
        }

        this.capacity = capacity;
        this.data = new HashMap<>();

        // Sentinel head and tail nodes
        this.head = new Node(null, null);
        this.tail = new Node(null, null);

        this.head.next = tail;
        this.tail.prev = head;
    }

    public Integer get(Integer key) {
        if (!data.containsKey(key)) {
            return -1;
        }

        moveToFront(key);

        Node node = data.get(key);
        return node.val;
    }

    public void put(Integer key, Integer val) {
        if (!data.containsKey(key)) {
            Node node = new Node(key, val);
            data.put(key, node);
            insertHead(node);
        } else {
            Node node = data.get(key);
            node.val = val;

            removeNode(node);
            insertHead(node);
        }

        if (data.size() > capacity) {
            evict();
        }
    }

    private void evict() {
        // Evict the previous node of the sentinel tail.
        Node node = tail.prev;

        removeNode(node);
        data.remove(node.key);
    }

    private void moveToFront(Integer key) {
        Node node = data.get(key);

        removeNode(node);
        insertHead(node);
    }

    private void removeNode(Node node) {
        Node prev = node.prev;
        Node next = node.next;

        prev.next = next;
        next.prev = prev;
    }

    private void insertHead(Node node) {
        Node next = head.next;

        // Insert the node next to the sentinel head.
        head.next = node;
        node.prev = head;
        node.next = next;
        next.prev = node;
    }
}

The time and space complexity remain unchanged from the previous implementation. By maintaining sentinel head and tail nodes, the implementation becomes simpler to understand.

Follow-up

Thread safety

The current implementation of the LRUCache can result in corrupted state if multiple threads put or get items simultaneously. For instance, the LRUCache manipulates the next and prev references of the nodes in many places to move the nodes to the head or evict a node. If multiple threads were to modify those references at the same time, the linked list could become invalid.

We can make the cache thread safe with a lock. Let us use a read/write lock to allow more than one threads to read the keys at the same time.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

class LRUCache {
    // ...
    private ReadWriteLock lock;

    public LRUCache(int capacity) {
        // ...
        this.lock = new ReentrantReadWritelock();
    }

    public Integer get(Integer key) {
        lock.readLock().lock();
        if (!data.containsKey(key)) {
            return -1;
        }
        lock.readLock().unlock();

        lock.writeLock().lock();
        moveToFront(key);
        Node node = data.get(key);
        lock.writeLock().unlock();

        return node.val;
    }

    public void put(Integer key, Integer val) {
        lock.writeLock().lock();
        // ... critical section
        lock.writeLock().unlock();
    }
}

As an improvement, we could use try-finally blocks to ensure that locks are released, and thereby preventing deadlocks.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

class LRUCache {
    // ...
    private ReadWriteLock lock;

    public LRUCache(int capacity) {
        // ...
        this.lock = new ReentrantReadWritelock();
    }

    public Integer get(Integer key) {
        lock.readLock().lock();
        try {
            if (!data.containsKey(key)) {
                return -1;
            }
        } finally {
            lock.readLock().unlock();
        }

        lock.writeLock().lock();
        try {
            moveToFront(key);
            Node node = data.get(key);
        } finally {
            lock.writeLock().unlock();
        }

        return node.val;
    }

    public void put(Integer key, Integer val) {
        lock.writeLock().lock();
        try {
            // ... critical section
        } finally {
            lock.writeLock().unlock();
        }
    }
}

While the critical sections currently do not throw any exceptions, it is a good practice to program defensively to avoid possibilities of a deadlock.

Lessons Learned

  1. Before implementing a data structure, clarify the APIs, data types, constraints, and edge cases.
  2. Break the logic down into maintainable pieces to make the data structure easier to understand and implement.
  3. When manipulating a linked list, try to create sentinel nodes to simplify the edge cases.
  4. A doubly-linked list can be used to maintain some changing order in a constant time.
  5. When designing a data structure, consider thread-safety as an edge case, and possibly cater for it using locks.