Skip to content

Commit 015842b

Browse files
committed
feat: Implement LRU Cache algorithm
1 parent 788d95b commit 015842b

File tree

1 file changed

+88
-0
lines changed

1 file changed

+88
-0
lines changed

data_structures/lru_cache.py

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
from __future__ import annotations
2+
3+
class Node:
4+
"""A node in the doubly linked list."""
5+
def __init__(self, key: int, val: int) -> None:
6+
self.key, self.val = key, val
7+
self.prev: Node | None = None
8+
self.next: Node | None = None
9+
10+
class LRUCache:
11+
"""
12+
A Least Recently Used (LRU) Cache data structure.
13+
>>> cache = LRUCache(2)
14+
>>> cache.put(1, 1)
15+
>>> cache.put(2, 2)
16+
>>> cache.get(1)
17+
1
18+
>>> cache.put(3, 3) # evicts key 2
19+
>>> cache.get(2)
20+
-1
21+
>>> cache.put(4, 4) # evicts key 1
22+
>>> cache.get(1)
23+
-1
24+
>>> cache.get(3)
25+
3
26+
>>> cache.get(4)
27+
4
28+
"""
29+
def __init__(self, capacity: int) -> None:
30+
if capacity <= 0:
31+
raise ValueError("Capacity must be a positive integer.")
32+
self.capacity = capacity
33+
self.cache: dict[int, Node] = {} # Maps key to node
34+
35+
self.head = Node(-1, -1)
36+
self.tail = Node(-1, -1)
37+
self.head.next = self.tail
38+
self.tail.prev = self.head
39+
40+
def _remove(self, node: Node) -> None:
41+
"""Helper to remove a node from the list."""
42+
if node.prev and node.next:
43+
prev, nxt = node.prev, node.next
44+
prev.next = nxt
45+
nxt.prev = prev
46+
47+
def _add(self, node: Node) -> None:
48+
"""Helper to add a node to the front of the list (most recent)."""
49+
if self.head.next:
50+
node.prev = self.head
51+
node.next = self.head.next
52+
self.head.next.prev = node
53+
self.head.next = node
54+
55+
def get(self, key: int) -> int:
56+
if key in self.cache:
57+
node = self.cache[key]
58+
self._remove(node)
59+
self._add(node)
60+
return node.val
61+
return -1
62+
63+
def put(self, key: int, value: int) -> None:
64+
"""
65+
Adds or updates a key-value pair in the cache.
66+
>>> cache = LRUCache(1)
67+
>>> cache.put(1, 10)
68+
>>> cache.get(1)
69+
10
70+
"""
71+
if key in self.cache:
72+
node = self.cache[key]
73+
node.val = value
74+
self._remove(node)
75+
self._add(node)
76+
else:
77+
new_node = Node(key, value)
78+
self.cache[key] = new_node
79+
self._add(new_node)
80+
81+
if len(self.cache) > self.capacity and self.tail.prev:
82+
lru = self.tail.prev
83+
self._remove(lru)
84+
del self.cache[lru.key]
85+
86+
if __name__ == "__main__":
87+
import doctest
88+
doctest.testmod()

0 commit comments

Comments
 (0)