LFU Cache [LeetCode 460]

LFUCache.py

# Reference : https://www.youtube.com/watch?v=Jn4mbZVkeik


from collections import OrderedDict

# A class to hold the value and the counter of the key
class KeyNode():
    def __init__(self, val, cnt):
        self.val = val
        self.cnt = cnt
        return

class LFUCache:

    def __init__(self, capacity: int):
        self.capacity = capacity
        self.cache = dict()
        self.counts = defaultdict(OrderedDict) # A dictionary of ordered list to get LFU key
        self.min_count = 0 # A variable to handle the current minimum count
        return

    def update(self, key: int) -> bool:
        
        # If key is present
        if key in self.cache:
            
            # Remove key from the current bucket
            self.counts[self.cache[key].cnt].pop(key, None)

            # Update the count of the key
            self.cache[key].cnt += 1

            # Add key to the new bucket
            self.counts[self.cache[key].cnt][key] = self.cache[key]
            
            # Update the min count if the current bucket is empty
            if not self.counts[self.min_count]:
                self.min_count += 1
            
            return True
        
        return False
    
    def get(self, key: int) -> int:
        
        # Update key count
        if self.update(key):
            return self.cache[key].val
        
        return -1
        
    def check(self) -> None:
        # If cache is full, remove LFU key
        if len(self.cache) == self.capacity:
            lfu_key, _ = self.counts[self.min_count].popitem(last=False)
            self.cache.pop(lfu_key, None)
        return
    
    def put(self, key: int, value: int) -> None:
        
        # Edge Case, if capacity is zero, do nothing
        if not self.capacity:
            return
       
        # Update the key count and value if key is present
        if self.update(key):
            self.cache[key].val = value
            
        else: # else if key is not present
            
            # Check if cache is not full, else do the needful
            self.check()
            
            # Add new key to cache, and to counts
            self.counts[1][key] = self.cache[key] = KeyNode(value, 1)
            self.min_count = 1
        
        return 


# Your LFUCache object will be instantiated and called as such:
# obj = LFUCache(capacity)
# param_1 = obj.get(key)
# obj.put(key,value)