Leetcode 432. All O`one Data Structure

class FrequencyNode:
    """Represents a frequency bucket containing all keys with a specific count."""
    def __init__(self, freq):
        self.freq = freq
        self.keys = set()  # Keys with this frequency
        self.prev = None
        self.next = None

class FrequencyTracker:
    def __init__(self):
        # Map from key to its current frequency
        self.key_freq = {}
        
        # Map from frequency to the FrequencyNode
        self.freq_nodes = {}
        
        # Dummy head and tail for doubly linked list of frequency buckets
        self.head = FrequencyNode(0)
        self.tail = FrequencyNode(float('inf'))
        self.head.next = self.tail
        self.tail.prev = self.head
    
    def _insert_node_after(self, new_node, prev_node):
        """Insert new_node after prev_node in the doubly linked list."""
        next_node = prev_node.next
        prev_node.next = new_node
        new_node.prev = prev_node
        new_node.next = next_node
        next_node.prev = new_node
    
    def _remove_node(self, node):
        """Remove node from the doubly linked list."""
        prev_node = node.prev
        next_node = node.next
        prev_node.next = next_node
        next_node.prev = prev_node
    
    def inc(self, key: str) -> None:
        # Get current frequency (0 if key doesn't exist)
        curr_freq = self.key_freq.get(key, 0)
        new_freq = curr_freq + 1
        
        # Update key frequency map
        self.key_freq[key] = new_freq
        
        # Remove key from old frequency bucket if it exists
        if curr_freq > 0:
            curr_node = self.freq_nodes[curr_freq]
            curr_node.keys.remove(key)
            
            # If bucket becomes empty, remove it
            if not curr_node.keys:
                self._remove_node(curr_node)
                del self.freq_nodes[curr_freq]
        
        # Find or create the new frequency bucket
        if new_freq not in self.freq_nodes:
            new_node = FrequencyNode(new_freq)
            self.freq_nodes[new_freq] = new_node
            
            # Insert new node in sorted position
            # Find the node to insert after
            if curr_freq > 0 and curr_freq in self.freq_nodes:
                insert_after = self.freq_nodes[curr_freq]
            else:
                # Find the correct position
                insert_after = self.head
                while insert_after.next.freq < new_freq:
                    insert_after = insert_after.next
            
            self._insert_node_after(new_node, insert_after)
        
        # Add key to new frequency bucket
        self.freq_nodes[new_freq].keys.add(key)
    
    def dec(self, key: str) -> None:
        # If key doesn't exist, do nothing
        if key not in self.key_freq:
            return
        
        curr_freq = self.key_freq[key]
        
        # Remove key from current frequency bucket
        curr_node = self.freq_nodes[curr_freq]
        curr_node.keys.remove(key)
        
        # If bucket becomes empty, remove it
        if not curr_node.keys:
            self._remove_node(curr_node)
            del self.freq_nodes[curr_freq]
        
        # Handle new frequency
        new_freq = curr_freq - 1
        
        if new_freq == 0:
            # Remove key completely
            del self.key_freq[key]
        else:
            # Update key frequency
            self.key_freq[key] = new_freq
            
            # Find or create new frequency bucket
            if new_freq not in self.freq_nodes:
                new_node = FrequencyNode(new_freq)
                self.freq_nodes[new_freq] = new_node
                
                # Insert before current position (since new_freq < curr_freq)
                insert_after = curr_node.prev if curr_freq in self.freq_nodes else self.head
                while insert_after.next.freq < new_freq:
                    insert_after = insert_after.next
                
                self._insert_node_after(new_node, insert_after)
            
            # Add key to new frequency bucket
            self.freq_nodes[new_freq].keys.add(key)
    
    def getMaxKey(self) -> str:
        # Max frequency node is just before tail
        if self.tail.prev == self.head:
            return ""
        
        max_node = self.tail.prev
        # Return any key from the max frequency bucket
        return next(iter(max_node.keys))
    
    def getMinKey(self) -> str:
        # Min frequency node is just after head
        if self.head.next == self.tail:
            return ""
        
        min_node = self.head.next
        # Return any key from the min frequency bucket
        return next(iter(min_node.keys))