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
