lfu_cache.py

import collections


class LFUCache(object):
    def __init__(self, n):
        self.capacity = n
        self.key_to_node = dict()
        self.count_to_nodes = collections.defaultdict(dict)
        self.min_count = 0

    class Node(object):
        def __init__(self, key, value):
            self.key = key
            self.value = value
            self.count = 1

        def __repr__(self):
            return f'Node({self.key}={self.value}, {self.count})'

    def set(self, key, value):
        if key in self.key_to_node:
            node = self.key_to_node[key]
            node.value = value
            self._increment_reference_count(node)
        else:
            if len(self.key_to_node) == self.capacity:
                self._evict()
            node = self.Node(key, value)
            self.key_to_node[key] = node
            self.min_count = 1
            count_dict = self.count_to_nodes[1]
            count_dict[key] = node

    def get(self, key):
        if key not in self.key_to_node:
            return None
        node = self.key_to_node[key]
        self._increment_reference_count(node)
        return node.value

    def _increment_reference_count(self, node):
        count = node.count
        count_dict = self.count_to_nodes[count]
        node.count += 1
        del(count_dict[node.key])
        next_count_dict = self.count_to_nodes[node.count]
        next_count_dict[node.key] = node
        if count == self.min_count and not count_dict:
            self.min_count = node.count

    def _evict(self):
        count_dict = self.count_to_nodes[self.min_count]
        key, node = next(iter(count_dict.items()))
        del(count_dict[key])
        del(self.key_to_node[key])