460. LFU Cache

/*
Time Complexity of all operations is O(1)
Space Complexity: O(capacity)
*/
class LFUCache {
    HashMap<Integer, Integer> keyValueMap;
    HashMap<Integer, Integer> keyCountMap;
    HashMap<Integer, LinkedHashSet<Integer>> countKeysMap;
    int maxCapacity;
    int curMin;

    public LFUCache(int capacity) {
        keyValueMap = new HashMap<Integer, Integer>();
        keyCountMap = new HashMap<Integer, Integer>();
        countKeysMap = new HashMap<Integer, LinkedHashSet<Integer>>();
        maxCapacity = capacity;
        curMin = -1;
    }

    public int get(int key) {
        if (!keyValueMap.containsKey(key)) {
            return -1;
        }
        int curCount = keyCountMap.get(key);
        keyCountMap.put(key, curCount+1);
        countKeysMap.get(curCount).remove(key);
        if (curMin == curCount && countKeysMap.get(curCount).isEmpty()) {
            curMin++;
            countKeysMap.remove(curCount);
        }
        addKeyToCountKeysMap(curCount+1, key);
        return keyValueMap.get(key);
    }

    public void put(int key, int value) {
        if (maxCapacity <= 0) {
            return;
        }
        if (keyValueMap.containsKey(key)) {
            keyValueMap.put(key, value);
            get(key);
            return;
        }

        if (keyValueMap.size() == maxCapacity) {
            int toBeRemoved = countKeysMap.get(curMin).iterator().next();
            countKeysMap.get(curMin).remove(toBeRemoved);
            if (countKeysMap.get(curMin).isEmpty()) {
                countKeysMap.remove(curMin);
            }
            keyValueMap.remove(toBeRemoved);
            keyCountMap.remove(toBeRemoved);
        }
        keyValueMap.put(key, value);
        keyCountMap.put(key, 1);
        curMin = 1;
        addKeyToCountKeysMap(1, key);
    }

    private void addKeyToCountKeysMap(int count, int key) {
        if (!countKeysMap.containsKey(count)) {
            countKeysMap.put(count, new LinkedHashSet<Integer>());
        }
        countKeysMap.get(count).add(key);
    }
}

/**
 * Your LFUCache object will be instantiated and called as such:
 * LFUCache obj = new LFUCache(capacity);
 * int param_1 = obj.get(key);
 * obj.put(key,value);
 */