cache_lfu.py

from typing import TypeAlias, Tuple, Hashable, TypeVar, Generic


FrequencyT: TypeAlias = int
ValueT = TypeVar("ValueT")
ValueFrequencyT: TypeAlias = tuple[ValueT, FrequencyT]
KeyT = TypeVar("KeyT", bound=Hashable)
KeysT: TypeAlias = set[KeyT]


class LFUCache(Generic[KeyT, ValueT]):
    def __init__(self, capacity: int):
        assert capacity > 0
        self.capacity: int = capacity
        self.cache: dict[KeyT, ValueFrequencyT[ValueT]] = {}
        self.freq: dict[FrequencyT, KeysT[KeyT]] = {}
        self.minFrequency: int = 0

    def get(self, key: KeyT) -> ValueT | None:
        if key not in self.cache:
            return None

        value, freq = self.cache[key]
        self._update(key, value, freq)
        return value

    def put(self, key: KeyT, value: ValueT) -> None:
        if key in self.cache:
            _, freq = self.cache[key]
            self._update(key, value, freq)
        else:
            if len(self.cache) >= self.capacity:
                self._evict()
            self.cache[key] = (value, 1)
            self.freq[1] = self.freq.get(1, set()) | {key}
            self.minFrequency = 1

    def _update(self, key: KeyT, value: ValueT, freq: FrequencyT) -> None:
        self.freq[freq].remove(key)
        if not self.freq[freq]:
            del self.freq[freq]
            if self.minFrequency == freq:
                self.minFrequency += 1

        self.cache[key] = (value, freq + 1)
        self.freq[freq + 1] = self.freq.get(freq + 1, set()) | {key}

    def _evict(self):
        key = self.freq[self.minFrequency].pop()
        if not self.freq[self.minFrequency]:
            del self.freq[self.minFrequency]
        del self.cache[key]


if __name__ == '__main__':
    lfu = LFUCache[int, str](3)
    lfu.put(1, 'A')
    lfu.put(2, 'B')
    assert lfu.get(1) == 'A', "Test 1 Failed"
    lfu.put(3, 'C')
    assert lfu.get(2) == 'B', "Test 2 Failed"
    lfu.put(4, 'D')
    assert lfu.get(3) == None, "Test 3 Failed"  # 'A' was evicted
    assert lfu.get(1) == 'A', "Test 4 Failed"
    assert lfu.get(2) == 'B', "Test 5 Failed"
    assert lfu.get(4) == 'D', "Test 6 Failed"
    print("OK")