Java | Just one TreeSet | add O(logM) avg O(1)
Last m stream elements are kept sorted using TreeSet. First K and Last K elements are identified by remembering the edge elements, and updated with each additio
Last m stream elements are kept sorted using TreeSet. First K and Last K elements are identified by remembering the edge elements, and updated with each addition (with O(logM)). Sum of m-2k "middle" elements is updated with each addition (and used to calculate average in O(1)). \nDuplicate elements are distinguished using a "timestamp" (here we can use just a simple serial integer, since number of calls to addElement <= 10^5)\n\nMore comments inline.\n```\nclass MKAverage {\n    private TreeSet<Pair<Integer, Integer>> order= new TreeSet<>((a,b)->{\n        int diff= a.getKey()-b.getKey();\n        return diff!=0 ? diff : a.getValue()-b.getValue();\n    });\n    private LinkedList<Pair<Integer, Integer>> stream= new LinkedList<>();\n    private Pair<Integer, Integer> kLeft= null, kRight= null;\n    private int m, k;\n    private int serial= 0;  // "timestamp"\n    private int midSum= 0, midSize= 0;\n\n    public MKAverage(int m, int k) {\n        this.m= m;\n        this.k= k;\n        this.midSize= m-2*k;\n    }\n    \n    private void init(){\n        // one off once we reach the size of m\n        // find kLeft and kRight elements (NOT in midSum)\n        // calculate initial midSum\n        int i=0;\n        for(Pair<Integer,Integer> elt:order){\n            if(i<k-1);\n            else if(i==k-1) kLeft= elt;\n            else if(i<m-k) midSum+= elt.getKey();\n            else if(i==m-k){kRight= elt; break;}\n            i++;\n        }\n    }\n    \n    public void addElement(int num) {\n        Pair<Integer, Integer> addElt= new Pair<>(num, serial++);\n        stream.addLast(addElt);\n        add(addElt);\n        if(stream.size()==m)\n            init();\n        else if(stream.size() > m)\n            remove(stream.removeFirst());\n    }\n    \n    private void add(Pair<Integer, Integer> elt){\n        if(order.size()<m){\n            order.add(elt);            \n            return;\n        }\n        // compare new elt with kLeft and kRight\n        int eltKey= elt.getKey(), kLeftKey= kLeft.getKey(), kRightKey= kRight.getKey();\n        order.add(elt);\n        if(eltKey < kLeftKey){\n            // add to left k and move kLeftKey to previous, \n            // adding kLeft to midSum\n            midSum+= kLeftKey;\n            // finding the elt immediately before the current kLeft\n            kLeft= order.floor(new Pair(kLeftKey, kLeft.getValue()-1));\n        }else if(eltKey < kRightKey){\n            // add to midSum, no change to kLeft or kRight\n            midSum+= eltKey;\n        }else{\n            // add to right k and move kRightKey to next\n            // adding kRight to midSum\n            midSum+= kRightKey;\n            // finding the elt immediately after the current kRight\n            kRight= order.ceiling(new Pair(kRightKey, kRight.getValue()+1));\n        }\n    }\n    \n    private void remove(Pair<Integer, Integer> elt){\n        // compare the elt being removed with kLeft and kRight\n        int eltKey= elt.getKey(), kLeftKey= kLeft.getKey(), kRightKey= kRight.getKey();\n        order.remove(elt);\n        if(eltKey <= kLeftKey){\n            // remove elt from left k and move kLeftKey to next, \n            // finding the elt immediately after the current kLeft\n            kLeft= order.ceiling(new Pair(kLeftKey, kLeft.getValue()+1));\n            // removing the new kLeft from midSum\n            midSum-= kLeft.getKey();\n        }else if(eltKey < kRightKey){\n            // remove from midSum, no change to kLeft or kRight\n            midSum-= eltKey;\n        }else{\n            // remove from right k and move kRightKey to previous\n            // finding the elt immediately before the current kRight\n            kRight= order.floor(new Pair(kRightKey, kRight.getValue()-1));\n            // removing the new kRight from  midSum\n            midSum-= kRight.getKey();\n        }\n    }    \n    \n    public int calculateMKAverage() {\n        return stream.size()==m ? midSum/midSize : -1;\n    }\n}
[]