example.py

from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, Optional, Generic, TypeVar

K = TypeVar("K")
V = TypeVar("V")


@dataclass
class _Node(Generic[K, V]):
    key: K
    value: V
    prev: Optional["_Node[K, V]"] = None
    next: Optional["_Node[K, V]"] = None


class LRUCache(Generic[K, V]):
    """
    LRU Cache: O(1) average get/put using:
      - dict: key -> node
      - doubly linked list: most-recent at front, least-recent at back
    """

    def __init__(self, capacity: int):
        if capacity <= 0:
            raise ValueError("capacity must be > 0")
        self._capacity = capacity
        self._map: Dict[K, _Node[K, V]] = {}

        # Sentinels to avoid edge cases.
        self._head = _Node(key=None, value=None)  # type: ignore[arg-type]
        self._tail = _Node(key=None, value=None)  # type: ignore[arg-type]
        self._head.next = self._tail
        self._tail.prev = self._head

    def get(self, key: K) -> Optional[V]:
        node = self._map.get(key)
        if node is None:
            return None
        self._move_to_front(node)
        return node.value

    def put(self, key: K, value: V) -> None:
        node = self._map.get(key)

        if node is not None:
            node.value = value
            self._move_to_front(node)
            return

        if len(self._map) >= self._capacity:
            lru = self._pop_lru()
            del self._map[lru.key]

        new_node = _Node(key=key, value=value)
        self._map[key] = new_node
        self._add_to_front(new_node)

    def __len__(self) -> int:
        return len(self._map)

    # ---------- linked list internals (O(1)) ----------

    def _add_to_front(self, node: _Node[K, V]) -> None:
        first = self._head.next
        node.prev = self._head
        node.next = first
        self._head.next = node
        if first is not None:
            first.prev = node

    def _remove(self, node: _Node[K, V]) -> None:
        prev = node.prev
        nxt = node.next
        if prev is not None:
            prev.next = nxt
        if nxt is not None:
            nxt.prev = prev
        node.prev = node.next = None

    def _move_to_front(self, node: _Node[K, V]) -> None:
        self._remove(node)
        self._add_to_front(node)

    def _pop_lru(self) -> _Node[K, V]:
        # Least-recently used is right before tail.
        lru = self._tail.prev
        if lru is None or lru is self._head:
            raise RuntimeError("LRU list is empty")
        self._remove(lru)
        return lru


# --------- quick example ----------
if __name__ == "__main__":
    cache = LRUCache[int, str](capacity=2)
    cache.put(1, "a")
    cache.put(2, "b")
    print(cache.get(1))  # "a" (1 becomes most-recent)
    cache.put(3, "c")    # evicts key 2
    print(cache.get(2))  # None