leetcode 1825. Finding MK Average

https://leetcode.com/problems/finding-mk-average/

You are given two integers, m and k, and a stream of integers. You are tasked to implement a data structure that calculates the MKAverage for the stream.

The MKAverage can be calculated using these steps:

1. If the number of the elements in the stream is less than m you should consider the MKAverage to be -1. Otherwise, copy the last m elements of the stream to a separate container.
2. Remove the smallest k elements and the largest k elements from the container.
3. Calculate the average value for the rest of the elements rounded down to the nearest integer.

Implement the MKAverage class:

• MKAverage(int m, int k) Initializes the MKAverage object with an empty stream and the two integers m and k.
• void addElement(int num) Inserts a new element num into the stream.
• int calculateMKAverage() Calculates and returns the MKAverage for the current stream rounded down to the nearest integer.

 

Example 1:

Input
["MKAverage", "addElement", "addElement", "calculateMKAverage", "addElement", "calculateMKAverage", "addElement", "addElement", "addElement", "calculateMKAverage"]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
Output
[null, null, null, -1, null, 3, null, null, null, 5]

Explanation
MKAverage obj = new MKAverage(3, 1); 
obj.addElement(3);        // current elements are [3]
obj.addElement(1);        // current elements are [3,1]
obj.calculateMKAverage(); // return -1, because m = 3 and only 2 elements exist.
obj.addElement(10);       // current elements are [3,1,10]
obj.calculateMKAverage(); // The last 3 elements are [3,1,10].
                          // After removing smallest and largest 1 element the container will be [3].
                          // The average of [3] equals 3/1 = 3, return 3
obj.addElement(5);        // current elements are [3,1,10,5]
obj.addElement(5);        // current elements are [3,1,10,5,5]
obj.addElement(5);        // current elements are [3,1,10,5,5,5]
obj.calculateMKAverage(); // The last 3 elements are [5,5,5].
                          // After removing smallest and largest 1 element the container will be [5].
                          // The average of [5] equals 5/1 = 5, return 5

 

Constraints:

• 3 <= m <= 105
• 1 <= k*2 < m
• 1 <= num <= 105
• At most 105 calls will be made to addElement and calculateMKAverage.

 

第一感是3个TreeMap。存前k大，前k小，和中间的数。但感觉应该有更好的方法。根据题意需要找到前k大的数，又需要求区间和，就自然想到线段树.写起来较容易出错。维护2个线段树数组，一个记录数的各数，一个记录区间值，注意一般线段树中[s，e]指固定的区间，这里类似线段数求第k小的数，所以[s,e]指第s小的值到第e小的值的区间。

另外第一次看到别人 二分+树状数组也能求前k大的值。https://leetcode.com/problems/finding-mk-average/discuss/1152431/Java-Fenwick-Tree-+-BinarySearch

 Segment Tree（68ms）

class MKAverage {
    LinkedList<Integer> queue;
    int[] count;
    long[] sum;
    int m,k;
    public MKAverage(int m, int k) {
        queue=new LinkedList<>();
        count=new int[400001];
        sum=new long[400001];
        this.m=m;
        this.k=k;
    }
    
    public void addElement(int num) {
        if(queue.size()==m){
            int v=queue.pollFirst();
            insert(1,0,100000,v,-1);
        }
        insert(1,0,100000,num,1);
        queue.addLast(num);
    }
    
    public int calculateMKAverage() {
        if(queue.size()<m)return -1;
        int s=k+1,e=m-k;
        return (int)(query(1,0,100000,s,e)/(m-2*k));
    }
    
    void insert(int node,int l,int r,int v,long d){
        count[node]+=d;
        sum[node]+=d*v;
        if(l==r){
            return;
        }
        int m=(l+r)/2;
        if(v<=m){
            insert(node<<1,l,m,v,d);        
        }else{
            insert(node<<1|1,m+1,r,v,d);     
        }
    }
    
    long query(int node,int l,int r,int s,int e){//线段中第s个到第e个
        if(l==r){//起始和结束最多出现2次此情况
            int c=e-s+1;
            return (long)c*l;
        }else if(count[node]==e-s+1){
            return sum[node];
        }else{
            int m=(l+r)/2;
            int c1=count[node<<1];
            int c2=count[node<<1|1];
            if(c1>=e){
                return query(node<<1,l,m,s,e);
            }else if(c1>=s){
                return query(node<<1,l,m,s,c1)+query(node<<1|1,m+1,r,1,e-c1);
            }else{//c1<s
                return query(node<<1|1,m+1,r,s-c1,e-c1);
            }
        }
    }
}

/**
 * Your MKAverage object will be instantiated and called as such:
 * MKAverage obj = new MKAverage(m, k);
 * obj.addElement(num);
 * int param_2 = obj.calculateMKAverage();
 */