3 TreeMaps (Clean Code)
Intuition\nThe idea is to have 3 treemaps tl, tm and th. For lowest k (tl), highest k (th) and all others (tm). The trick is we need to balance the maps when an
# Intuition\nThe idea is to have 3 treemaps **tl**, **tm** and **th**. For lowest k (tl), highest k (th) and all others (tm). The trick is we need to balance the maps when an element is added/removed. \n\n\n# Code\n```\nclass MKAverage {\n    int mx = (int)2e5, mn = 0, i = 0, m, k;\n    long sm = 0l;\n    TreeMap<Integer, Integer> tl = new TreeMap(), tm = new TreeMap(), th = new TreeMap(); \n    HashMap<Integer, Integer> ind = new HashMap();\n\n    public MKAverage(int m, int k) {\n        this.m = m;\n        this.k = k;\n    }\n    \n    public void addElement(int num) {\n        addNewNum(num);\n        if(i > m) removeOldNum();\n    }\n    \n    public int calculateMKAverage() {\n        return i >= m ? (int)(sm / ((m - 2 * k)*1l)) : -1;\n    }\n\n    void addNewNum(int num){\n        ind.put(i++, num);\n        add(tl, num);\n        if(i <= k) return;\n        int v1 = removeHighest(tl);\n        \n        add(tm, v1);\n        int v2 = removeHighest(tm);\n        sm += v1 - v2;\n\n        add(th, v2);\n        if(i <= 2*k) return;\n        int v3 = removeLowest(th);\n\n        add(tm, v3);\n        sm += v3;\n    }\n\n    void removeOldNum() {\n        int prev = ind.get(i - m - 1);\n\n        if(tl.get(prev) != null){\n            remove(tl, prev);\n            int low = removeLowest(tm);\n            sm -= low;\n            add(tl, low);\n        } else if (th.get(prev) != null){\n            remove(th, prev);\n            int high = removeHighest(tm);\n            sm -= high;\n            add(th, high);\n        } else {\n            remove(tm, prev);\n            sm -= prev;\n        }\n\n        ind.remove(i - m - 1);\n    }\n\n    void remove(Map<Integer, Integer> tree, int val) {\n        tree.put(val, tree.getOrDefault(val, 0) - 1);\n        if(tree.get(val) == 0) tree.remove(val);\n    }\n\n    void add(Map<Integer, Integer> tree, int val){\n        tree.put(val, tree.getOrDefault(val, 0) + 1);\n    }\n\n    int removeLowest(TreeMap<Integer, Integer> tree){\n        int low = tree.ceilingKey(mn);\n        tree.put(low, tree.get(low) - 1);\n        if(tree.get(low) == 0) tree.remove(low);\n        return low;\n    }\n\n    int removeHighest(TreeMap<Integer, Integer> tree){\n        int high = tree.floorKey(mx);\n        tree.put(high, tree.get(high) - 1);\n        if(tree.get(high) == 0) tree.remove(high);\n        return high;\n    }\n}\n```
['Java']