[LeetCode] 432. All O`one Data Structure

Problem

Implement a data structure supporting the following operations:

Inc(Key) - Inserts a new key with value 1. Or increments an existing key by 1. Key is guaranteed to be a non-empty string.
Dec(Key) - If Key's value is 1, remove it from the data structure. Otherwise decrements an existing key by 1. If the key does not exist, this function does nothing. Key is guaranteed to be a non-empty string.
GetMaxKey() - Returns one of the keys with maximal value. If no element exists, return an empty string "".
GetMinKey() - Returns one of the keys with minimal value. If no element exists, return an empty string "".
Challenge: Perform all these in O(1) time complexity.

Note

Use DoublyLinkedList to get max/min key (max/min count)
We put max count Bucket at tail, min count Bucket at head, O(1)

So the node(Bucket) must contain count property
And we can use map <count, node> to access each node in O(1)
Each Bucket could have multiple keys of same count, so use Set to store keys

To get count of a key, use Map <count, key> to achieve O(1)

Solution

class AllOne {
    private class Bucket {
        int count;
        Set<String> keySet;
        Bucket prev;
        Bucket next;
        public Bucket(int count) {
            this.count = count;
            this.keySet = new HashSet<>();
        }
    }

    private Bucket head;
    private Bucket tail;
    private Map<Integer, Bucket> countBucketMap;
    private Map<String, Integer> keyCountMap;
    
    public AllOne() {
        head = new Bucket(Integer.MIN_VALUE);
        tail = new Bucket(Integer.MAX_VALUE);
        head.next = tail;
        tail.prev = head;
        countBucketMap = new HashMap<>();
        keyCountMap = new HashMap<>();
    }
    
    public void inc(String key) {
        if (keyCountMap.containsKey(key)) {
            //existing key
            updateKey(key, 1);
        } else {
            //new key
            keyCountMap.put(key, 1);
            if (head.next.count != 1) {
                addBucketAfter(head, new Bucket(1));
            }
            head.next.keySet.add(key);
            countBucketMap.put(1, head.next);
        }
    }
    
    public void dec(String key) {
        if (keyCountMap.containsKey(key)) {
            int count = keyCountMap.get(key);
            if (count == 1) {
                keyCountMap.remove(key);
                removeKeyFromBucket(countBucketMap.get(count), key);
            } else {
                updateKey(key, -1);
            }
        }
    }
    
    private void updateKey(String key, int i) {
        int count = keyCountMap.get(key);
        keyCountMap.put(key, count+i);
        Bucket pre = countBucketMap.get(count);
        Bucket cur;
        if (countBucketMap.containsKey(count+i)) {
            cur = countBucketMap.get(count+i);
        } else {
            cur = new Bucket(count+i);
            countBucketMap.put(count+i, cur);
            if (i == 1) {
                addBucketAfter(pre, cur);
            } else {
                addBucketAfter(pre.prev, cur);
            }
        }
        cur.keySet.add(key);
        removeKeyFromBucket(pre, key);
    }
    
    private void addBucketAfter(Bucket pre, Bucket cur) {
        Bucket next = pre.next;
        pre.next = cur;
        cur.prev = pre;
        cur.next = next;
        next.prev = cur;
    }
    
    private void removeKeyFromBucket(Bucket bucket, String key) {
        bucket.keySet.remove(key);
        if (bucket.keySet.size() == 0) {
            bucket.prev.next = bucket.next;
            bucket.next.prev = bucket.prev;
            bucket.prev = null;
            bucket.next = null;
            countBucketMap.remove(bucket.count);
        }
    }
    
    public String getMaxKey() {
        return tail.prev == head ? "" : (String) tail.prev.keySet.iterator().next();
    }
    
    public String getMinKey() {
        return head.next == tail ? "" : (String) head.next.keySet.iterator().next();
    }
}

Reference: https://leetcode.com/problems...