432. All O`one Data Structure

class Bucket:
    def __init__(self,cnt = 0):
        self.cnt = cnt
        self.keySet = set()
        self.pre = None 
        self.nxt = None 
        
    def _get_any_key(self):
        if len(self.keySet)!=0:
            res = self.keySet.pop()
            self.keySet.add(res)
            return res 
        else:
            return None
        
# max is on the tail while min is on the head 
class AllOne:

    def __init__(self):
        self.head = Bucket() # head maintains max val, tail maintains min val 
        self.tail = Bucket()
        self.cnt_bucket_map = {}
        self.key_cnt_map = {}
        self.head.nxt = self.tail
        self.tail.pre = self.head
        
    def _change_key(self,key,offset):
        precnt = self.key_cnt_map[key]
        # update key_cnt map 
        self.key_cnt_map[key] =  precnt + offset
        cur_bucket = self.cnt_bucket_map[precnt]
        # update bucket 
        newcnt = precnt + offset
        
        # check if new bucket in bucket list 
        if newcnt in self.cnt_bucket_map:
            new_bucket = self.cnt_bucket_map[newcnt]
        else:
            new_bucket = Bucket(newcnt)
            # update cnt_bucket map 
            self.cnt_bucket_map[newcnt] = new_bucket
            if offset == 1: 
                self._insert_after(cur_bucket,new_bucket)
            else:
                self._insert_after(cur_bucket.pre,new_bucket)
        # add key to new bucket  
        new_bucket.keySet.add(key)
       # print(f"cnt is {new_bucket.cnt} set is {new_bucket.keySet}")
        # remove key from old bucket 
        self._remove_from_bucket(cur_bucket,key)
        #print(f" remove  {key} from bucket {cur_bucket.cnt}")
       # print(f" after remove  {key} from bucket {cur_bucket.cnt}, bucket {cur_bucket.cnt} contains {cur_bucket.keySet}")

    def inc(self, key: str) -> None:
        precnt = self.key_cnt_map.get(key,-1)
        if precnt != -1: #seen before 
            self._change_key(key,1)
        else: # not seen 
            self.key_cnt_map[key] = 1
            if self.head.nxt.cnt != 1: # min bucket bigger than 1 
                self._insert_after(self.head,Bucket(1))
            self.head.nxt.keySet.add(key)
            self.cnt_bucket_map[1] = self.head.nxt
        
    def dec(self, key: str) -> None:
        if key in self.key_cnt_map:
            cnt = self.key_cnt_map[key]
            bucket = self.cnt_bucket_map[cnt]
            if cnt == 1:
                self.key_cnt_map.pop(key)
                self._remove_from_bucket(bucket,key)
            else:
                self._change_key(key,-1)
            
        
    def getMaxKey(self) -> str:
        return self.tail.pre._get_any_key() if self.tail.pre != self.head else ''
        

    def getMinKey(self) -> str:
        #print(self.head.nxt.keySet)
        return self.head.nxt._get_any_key() if self.head.nxt != self.tail else ''
    
    
    def _remove_from_bucket(self,pre_bucket,key):
        pre_bucket.keySet.remove(key)
        if len(pre_bucket.keySet) == 0:
            self.cnt_bucket_map.pop(pre_bucket.cnt)
            self._remove_bucket_from_list(pre_bucket)
            
    def _remove_bucket_from_list(self,bucket):
        bucket.pre.nxt = bucket.nxt
        bucket.nxt.pre = bucket.pre
        bucket.pre = None
        bucket.nxt = None 
        
    def _insert_after(self,pre_node,new_node):
        new_node.pre = pre_node  # new_node.pre may run into NoneType object if 
        new_node.nxt = pre_node.nxt
        pre_node.nxt.pre = new_node #if new_node.nxt else None 
        pre_node.nxt = new_node
        
    # def _insert_before(self,pre_node,new_node):
    #     new_node.pre = pre_node.pre
    #     new_node.nxt = pre_node
    #     pre_node.pre.nxt = new_node
    #     pre_node.pre = new_node 
# Your AllOne object will be instantiated and called as such:
# obj = AllOne()
# obj.inc(key)
# obj.dec(key)
# param_3 = obj.getMaxKey()
# param_4 = obj.getMinKey()