[Java] 3 TreeSets + 1 Queue, clear and easy-to-write code
Similar problems\nThis problem is a variant of\n- 295. Find Median from Data Stream\n- 480. Sliding Window Median\n\nAnd I strongly suggest you read this post, 
## Similar problems\nThis problem is a variant of\n- [295. Find Median from Data Stream](https://leetcode.com/problems/find-median-from-data-stream/)\n- [480. Sliding Window Median](https://leetcode.com/problems/sliding-window-median/)\n\nAnd I strongly suggest you read [this post](https://leetcode.com/problems/sliding-window-median/discuss/96346/), which shows an elegant way to rebalance between disjoint sorted sets.\n\n## Intuition\n\nLet\'s get back to this problem. We\'re asked to maintain the last `m` elements of a stream, while return its average excluding top `k` and bottom `k` elements. It\'s easy to infer that some kind of sorted set data structure should be used, and logarithmic maintenance time is required.\n\n### Difficulties\n\nThe major difficulties include\n1. How to maintain a dynamic sorted set in the first place\n2. How to distinguish the top `k` and bottom `k` elements from the candidates for MK average\n\nFor 1, C++ users are lucky since there is built-in multiset. However, it\'s not difficult to implement one in Java. The key is to differentiate between elements with the same value, which can be done using a logical counter incrementing by 1, and this assigned unique identifier can be stored with the value. Several lines of code should suffice to define such a class and corresponding comparator (or alternatively the `Comparable` interface).\n\nFor 2, similar to the aforementioned problems, we can use multiple disjoint dynamic sorted sets to distinguish between and maintain subsets of a population. In particular, a top `k` set, a bottom `k` set and a window of candidates for evaluation of MK average should be identified in this problem. Thus 3 multisets are needed. In addition, we should be able to remove expired elements, which can be done by recording the entry order of elements with a single queue.\n\n###  Invariant Maintenance\n\nThe most complicated part is maintain the invariant that\n> Top `k` set and bottom `k` set contain the corresponding top or bottom `k` elements, while the remaining set contains the `m - 2k` elements left.\n\nAgain, I would recommend you reading the post above to learn about the intuition to rebalance between multiple sorted sets. The most important part is to define a **balancing operation**.\n\n## Implementation\n\n```java\nclass MKAverage {\n    static class Element {\n        int id;\n        int val;\n        \n        Element(int id, int val) {\n            this.id = id;\n            this.val = val;\n        }\n    }\n    \n    int m;\n    int k;\n    int size = 0;\n    int sum = 0;\n    Comparator<Element> cmp = Comparator.<Element>comparingInt(e -> e.val).thenComparingInt(e -> e.id);\n    TreeSet<Element> window = new TreeSet<>(cmp);\n    TreeSet<Element> topK = new TreeSet<>(cmp);\n    TreeSet<Element> bottomK = new TreeSet<>(cmp);\n    Queue<Element> queue = new ArrayDeque<Element>();\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        var element = new Element(size++, num);\n        queue.offer(element);\n        window.add(element);\n        sum += num;\n        \n        if (size == m) {\n            while (window.size() > m - 2 * k) {\n                balance();\n            }\n        } else if (size > m) {\n            if (window.last().val > topK.first().val) {\n                move(window, topK);\n            } else {\n                move(window, bottomK);\n            }  \n            var toRemove = queue.poll();\n            topK.remove(toRemove);\n            bottomK.remove(toRemove);\n            if (window.remove(toRemove)) {\n                sum -= toRemove.val;\n            }\n            balance();\n        }\n    }\n    \n    private void move(TreeSet<Element> from, TreeSet<Element> to) {\n        var element = (from == topK || to == bottomK) ? from.pollFirst() : from.pollLast();\n        to.add(element);\n        sum += to == window ? element.val : -element.val;\n    }\n    \n    private void balance() {\n        if (topK.size() > k) {\n            move(topK, window);\n        } else if (bottomK.size() > k) {\n            move(bottomK, window);\n        }\n        \n        if (topK.size() < k) {\n            move(window, topK);\n        } else if (bottomK.size() < k) {\n            move(window, bottomK);\n        }\n    }\n    \n    public int calculateMKAverage() {\n        return size < m ? -1 : sum / (m - 2 * k);\n    }\n}\n```
[]