99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
"""O(1) LRU cache: doubly linked list + hash map.
|
|
|
|
Design notes:
|
|
- Each node holds key, value, prev, next.
|
|
- `_map` maps key -> node.
|
|
- Head is MRU sentinel, tail is LRU sentinel. Most-recently-used sits right after head.
|
|
- `get` and `put`-on-existing must BOTH promote the node to MRU.
|
|
- `put`-on-new inserts at MRU; if over capacity, evicts the node before tail (LRU).
|
|
- Eviction must delete from `_map` (common bug: forget this).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
|
|
class _Node:
|
|
__slots__ = ("key", "value", "prev", "next")
|
|
|
|
def __init__(self, key: Any, value: Any) -> None:
|
|
self.key = key
|
|
self.value = value
|
|
self.prev: Optional["_Node"] = None
|
|
self.next: Optional["_Node"] = None
|
|
|
|
|
|
class LRUCache:
|
|
def __init__(self, capacity: int) -> None:
|
|
if not isinstance(capacity, int) or isinstance(capacity, bool):
|
|
raise TypeError("capacity must be an int")
|
|
if capacity <= 0:
|
|
raise ValueError("capacity must be positive")
|
|
self.capacity = capacity
|
|
self._map: dict[Any, _Node] = {}
|
|
# sentinels
|
|
self._head = _Node(None, None) # MRU side
|
|
self._tail = _Node(None, None) # LRU side
|
|
self._head.next = self._tail
|
|
self._tail.prev = self._head
|
|
|
|
# ---- internal list ops ----
|
|
def _detach(self, node: _Node) -> None:
|
|
p = node.prev
|
|
n = node.next
|
|
assert p is not None and n is not None
|
|
p.next = n
|
|
n.prev = p
|
|
node.prev = None
|
|
node.next = None
|
|
|
|
def _attach_mru(self, node: _Node) -> None:
|
|
# insert right after head
|
|
first = self._head.next
|
|
self._head.next = node
|
|
node.prev = self._head
|
|
node.next = first
|
|
assert first is not None
|
|
first.prev = node
|
|
|
|
def _promote(self, node: _Node) -> None:
|
|
self._detach(node)
|
|
self._attach_mru(node)
|
|
|
|
# ---- public API ----
|
|
def get(self, key: Any) -> Any:
|
|
node = self._map.get(key)
|
|
if node is None:
|
|
return -1
|
|
self._promote(node)
|
|
return node.value
|
|
|
|
def put(self, key: Any, value: Any) -> None:
|
|
node = self._map.get(key)
|
|
if node is not None:
|
|
# update value AND recency (subtle bug if you skip promote)
|
|
node.value = value
|
|
self._promote(node)
|
|
return
|
|
node = _Node(key, value)
|
|
self._map[key] = node
|
|
self._attach_mru(node)
|
|
if len(self._map) > self.capacity:
|
|
# evict LRU: node just before tail
|
|
lru = self._tail.prev
|
|
assert lru is not None and lru is not self._head
|
|
self._detach(lru)
|
|
del self._map[lru.key]
|
|
|
|
def size(self) -> int:
|
|
return len(self._map)
|
|
|
|
# helpers for tests / introspection
|
|
def order_mru_to_lru(self) -> list[Any]:
|
|
out = []
|
|
cur = self._head.next
|
|
while cur is not None and cur is not self._tail:
|
|
out.append(cur.key)
|
|
cur = cur.next
|
|
return out
|