72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
from collections import OrderedDict
|
|
from typing import DefaultDict, Dict, Generic, Iterator, Tuple, TypeVar
|
|
|
|
K = TypeVar("K")
|
|
V = TypeVar("V")
|
|
|
|
__all__ = ["LFUCache"]
|
|
|
|
|
|
class LFUCache(Generic[K, V]):
|
|
def __init__(self, capacity: int):
|
|
if capacity <= 0:
|
|
raise ValueError("Capacity must be positive")
|
|
|
|
self.capacity = capacity
|
|
self.cache: Dict[K, Tuple[V, int]] = {} # To store key-value pairs
|
|
self.freq_count: DefaultDict[int, OrderedDict[K, V]] = DefaultDict(
|
|
OrderedDict
|
|
) # To store frequency of each key
|
|
self.min_freq = 0 # To keep track of the minimum frequency
|
|
|
|
def get(self, key: K) -> V:
|
|
if key in self.cache:
|
|
value, freq = self.cache[key]
|
|
# Update frequency and move the key to the next frequency bucket
|
|
self.freq_count[freq].pop(key)
|
|
if not self.freq_count[freq]: # If the current frequency has no keys, update min_freq
|
|
del self.freq_count[freq]
|
|
if freq == self.min_freq:
|
|
self.min_freq += 1
|
|
freq += 1
|
|
self.freq_count[freq][key] = value
|
|
self.cache[key] = (value, freq)
|
|
return value
|
|
raise KeyError(f"Key {key} not found")
|
|
|
|
def put(self, key: K, value: V) -> None:
|
|
if key in self.cache:
|
|
_, freq = self.cache[key]
|
|
# Update frequency and move the key to the next frequency bucket
|
|
self.freq_count[freq].pop(key)
|
|
if not self.freq_count[freq]:
|
|
del self.freq_count[freq]
|
|
if freq == self.min_freq:
|
|
self.min_freq += 1
|
|
freq += 1
|
|
self.freq_count[freq][key] = value
|
|
self.cache[key] = (value, freq)
|
|
else:
|
|
# Check if cache is full, evict the least frequently used item
|
|
if len(self.cache) == self.capacity:
|
|
evicted_key, _ = self.freq_count[self.min_freq].popitem(last=False)
|
|
del self.cache[evicted_key]
|
|
|
|
# Add the new key-value pair with frequency 1
|
|
self.cache[key] = (value, 1)
|
|
self.freq_count[1][key] = value
|
|
self.min_freq = 1
|
|
|
|
def remove_key(self, key: K) -> None:
|
|
if key in self.cache:
|
|
_, freq = self.cache[key]
|
|
self.freq_count[freq].pop(key)
|
|
if not self.freq_count[freq]: # If the current frequency has no keys, update min_freq
|
|
del self.freq_count[freq]
|
|
if freq == self.min_freq:
|
|
self.min_freq += 1
|
|
del self.cache[key]
|
|
|
|
def __iter__(self) -> Iterator[K]:
|
|
yield from self.cache
|