[LeetCode] 求出 MK 平均值 树状数组+二分

题目链接：求出MK平均值

题面：

思路：

起初看到题目，觉得难以下手（我太菜了T^T），后面看到num最大只有1e5，便想到根据num的值域建立树状数组...详情请看代码。

代码：

class MKAverage {
    public BIT bit;   //维护数字的个数
    public BIT sum;   //维护数字的和
    public ArrayList<Integer> buf;
    public int m;
    public int k;

    public MKAverage(int m, int k) {
        buf=new ArrayList<Integer>();
        bit=new BIT(100010);
        sum=new BIT(100010);
        this.m=m;
        this.k=k;
    }
    
    public void addElement(int num) {
        if(buf.size()<m){
            buf.add(num);
            bit.add(num,1);
            sum.add(num,num);
        }else{
            bit.add(buf.get(buf.size()-m),-1);
            sum.add(buf.get(buf.size()-m),-buf.get(buf.size()-m));
            buf.add(num);
            bit.add(num,1);
            sum.add(num,num);
        }
    }
    
    public int calculateMKAverage() {
        if(buf.size()<m){
            return -1;
        }
        int l=0;                    //此处要注意不要写成l=1,r=100000,刚开始因为这个wa了一发,因为可能num全为1或者全为100000
        int r=100001;
        int lidx=0;
        while(l<=r){
            int mid=(l+r)>>1;
            if(bit.ask(1,mid)<=k){
                lidx=mid;
                l=mid+1;
            }else{
                r=mid-1;
            }
        }
        l=0;
        r=100001;
        int ridx=0;
        while(l<=r){
            int mid=(l+r)>>1;
            if(bit.ask(mid,100000)<=k){
                ridx=mid;
                r=mid-1;
            }else{
                l=mid+1;
            }
        }
        long cntl=bit.ask(1,lidx);
        long cntr=bit.ask(ridx,100000);
        long all=sum.ask(1,100000);
        long lsum=sum.ask(1,lidx);
        long rsum=sum.ask(ridx,100000);
        if(cntl!=k)
            lsum+=(lidx+1)*(k-cntl);
        if(cntr!=k)
            rsum+=(ridx-1)*(k-cntr);
        return (int)((all-lsum-rsum)/(m-2*k));
    }

    static class BIT{
		public long[] C;        //注意本题在求和的过程中会爆int 
		
		public BIT(int n) {
			C=new long[n+1];
		}
		
		public int lowBit(int x) {
			return x&(-x);
		}
		
		public void add(int x,long y) {
			for(int i=x;i<C.length;i+=lowBit(i)) {
				C[i]+=y;
			}
		}
		
		public long ask(int x) {
			long res=0;
			
			for(int i=x;i>0;i-=lowBit(i))
				res+=C[i];
			
			return res;
		}
		
		//C[i]表示i这个数有C[i]个，相当于权值BIT，查询第k大。
	    public int kth_element(int k) {
	        int p = 0;
	        for (int lim = 1 << 20; lim>0; lim /= 2)
	            if (p + lim < C.length && C[p + lim] < k) {
	                p += lim;
	                k -= C[p];
	            }
	        return p + 1;
	    }
	    
	    public void add(int l,int r,long val) {
	    	this.add(l, val);
	    	this.add(r+1, -val);
	    }
	    
	    public long ask(int l,int r) {
	    	return this.ask(r)-this.ask(l-1);
	    }
	}
}

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