Need some feedback on this implementation of LFU cache in python3.
Original problem : https://leetcode.com/problems/lfu-cache/
Would really appreciate some feedback on readability, understandability and areas of improvement.
""" ListNode class to create new nodes objects @variables: key - Key from key,value pair stored by node value - Value from key,value pair stored by node next - Link to next ListNode object of doubly linked list prev - Link to previous ListNode object of doubly linked list frequency - Holds the number of times this node was accessed in LFUCache.get() and LFUCache.put() methods. Default is 1 at node creation on LFUCache.put() call """ class ListNode: def __init__(self, key=0, val=0, next=None, prev=None): self.key = key self.val = val self.next = next self.prev = prev self.frequency = 1 """ Main class to create LFU cache. The idea is to maintain cache linked list in order of frequency and LRU within same frequency A hash map (self.hash_map) will hold node key as key and pointer to node as value. Another hash map will hold distinct node frequencies as key and pointer to head node for the frequency. ______ ______ ______ ______ ______ | | | | | | hash_map | 23 | 56 | 14 | 6 | 19 | |______|______|______|______|______| | | | | | _____________| ____| | |_____ |_______________ / / | \ \ _______ ___V___ ____V__ __V____ _V_____ ___V___ _______ | | <--- | | <--- | | <--- | | <--- | | <--- | | <--- | | LFUCache | Head | | 23(5) | | 56(3) | | 14(2) | | 6(2) | | 19(1) | | Tail | |_______| ---> |_______| ---> |_______| ---> |_______| ---> |_______| ---> |_______| ---> |_______| ^ ^ ^ ^ \ \______ | / \______________ \ | _____________________/ _\___ __\__ __|__ __/__ | | | | | freq_map | 5 | 3 | 2 | 1 | |_____|_____|_____|_____| """ class LFUCache: def __init__(self, capacity: int): self.capacity = capacity self.hash_map = {} # Hash map to hold cache keys as key and link to node as value self.freq_map = {} # Hash map to hold frequency as key and link to head node for that frequency # Dummy head and tail cache nodes to make opertions on edge cases(one node and two node cache) simple self.head = ListNode(0,0) self.tail = ListNode(0,0) # Initially set tail (dummy node) as head node for frequency 1 # so that first cache node is added before tail node self.freq_map[1] = self.tail # Link head cache node and tail cache nodes together self.head.next,self.tail.prev = self.tail,self.head """ This method will get the value of key (input) if it exists in the LFU cache else returns -1. If key exists then this method will call self._move() to move the node to front of it's new frequency queue """ def get(self, key: int) -> int: if self.hash_map.get(key): # Key exists, get link to node from hash map node = self.hash_map.get(key) # Since frequency changed, move this node to # front of new frequency queue self._move(node) return node.val else: return -1 """ This method will update the value of key (input) if it exists in the LFU cache else inserts new node into the appropriate position on LFU cache """ def put(self, key: int, value: int) -> None: # Handle edge case when capacity is 0 if self.capacity == 0: return if self.hash_map.get(key): # Key exists, get link to node from hash map self.hash_map.get(key).val = value # Move this node to front of new frequency queue self._move(self.hash_map.get(key)) else: # Key does not exist, need to create a new node. Check if capcaity is full if len(self.hash_map) == self.capacity: # If node to be removed is the head node for it's frequency then remove it from frequency map # If frequency of node to be removed is 1 then we need to reset head node for frequency 1 to tail node if self.freq_map.get(self.tail.prev.frequency) == self.tail.prev and self.tail.prev.frequency == 1: self.freq_map[1] = self.tail elif self.freq_map.get(self.tail.prev.frequency) == self.tail.prev: self.freq_map.pop(self.tail.prev.frequency) # Remove last node (tail.prev) from the cache list self._remove(self.tail.prev) # Create new node and # - Add it to cache list # - Set it as head node for frequency 1 node = ListNode(key,value) head = self.freq_map.get(1) self.freq_map[1] = node self._add(node, head) """ This method will add a new node(input) before the head(input) node on LFU cache """ def _add(self, node :ListNode, head :ListNode) -> None: # When adding in the middle of cache we don't have dummy head node # so we set head.prev to play as dummy head node if head != self.head: head = head.prev # Add node operation node.next = head.next head.next.prev = node head.next = node node.prev = head self.hash_map[node.key] = node """ This method will remove a node(input) from LFU cache and hash_map """ def _remove(self, node :ListNode) -> None: node.prev.next = node.next node.next.prev = node.prev node.next = None node.prev = None self.hash_map.pop(node.key) def _move(self, node :ListNode) -> None: new_frequency = node.frequency + 1 # We need to find head node, before which we need to place the input node # If new frequency already exists in the cache list then get the head node for that frequency # else if new frequency does not exist then node will remain at same place in cache so set head as node.next # else get the head node for current frequency of input node if self.freq_map.get(new_frequency): head = self.freq_map.get(new_frequency) elif self.freq_map.get(node.frequency) == node: head = node.next else: head = self.freq_map.get(node.frequency) # Set new head nodes for current frequency of input node if self.freq_map.get(node.frequency) == node and node.frequency == node.next.frequency: self.freq_map[node.frequency] = node.next elif self.freq_map.get(node.frequency) == node: self.freq_map.pop(node.frequency) self._remove(node) self._add(node,head) node.frequency = new_frequency # Set this node as head node for new frequency (current frequency + 1) self.freq_map[new_frequency] = node 
