topic link

Leetcode.1825 find the average value of MK

topic description

You are given two integers m and k, and a number of integers in the form of a data stream. You need to implement a data structure to calculate the MK average of this data stream.

The MK average is calculated as follows:

If there are less than m integers in the data stream, the MK average is -1, otherwise the last m elements of the data stream are copied into a separate container.
Remove the smallest k numbers and the largest k numbers from this container.
Calculates the average of the remaining elements, rounding down to the nearest integer.
Please implement the MKAverage class:

MKAverage(int m, int k) initializes an MKAverage object with an empty stream and two integers m and k.
void addElement(int num) inserts a new element num into the data stream.
int calculateMKAverage() calculates and returns the MK average for the current data stream, and the result needs to be rounded down to the nearest integer.

Example 1:

enter:
["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 element is [3]
obj.addElement(1); // current element is [3,1]
obj.calculateMKAverage(); // returns -1 because m = 3 but there are only 2 elements in the stream
obj.addElement(10); // current element is [3,1,10]
obj.calculateMKAverage(); // The last 3 elements are [3,1,10]
// After removing the smallest and largest 1 element, the container is [3]
// The mean of [3] is equal to 3/1 = 3, so returns 3
obj.addElement(5); // current element is [3,1,10,5]
obj.addElement(5); // current element is [3,1,10,5,5]
obj.addElement(5); // current element is [3,1,10,5,5,5]
obj.calculateMKAverage(); // The last 3 elements are [5,5,5]
// After deleting the smallest and largest 1 element, the container is [5]
// The average of [5] equals 5/1 = 5, so returns 5

hint:

• 3 < = m < = 1 0 5 3 <= m <= 10^5 3<=m<=105
• 1 < = k ∗ 2 < m 1 <= k*2 < m 1<=k∗2<m
• 1 < = n u m < = 1 0 5 1 <= num <= 10^5 1<=num<=105
• The total number of operations of addElement and calculateMKAverage should not exceed 1 0 5 10^5 105 times.

analyze:

This question needs to maintain a queue of size m, and the elements in it are the latest m elements. In addition, we also try to quickly find the average value of the sum of the remaining elements (rounded down) after removing the largest k elements and the smallest k elements from the elements in the queue.
We can maintain this data structure using two tree arrays and a queue.

Tree array can see this blog: tree array

• A queue q, in which only the latest m elements are stored
• A tree array s, storing the sum of inserted elements
• An array array cnt, storing the number of times to insert elements

code:

using LL = long long;
const int N = 1e5+10;

class MKAverage {
public:
    //Observe the data range, the s array of the record element sum may explode int, so use long long
    //For convenience, long long is used uniformly here
    LL cnt[N],s[N];
    
    int k,m;
    queue<int> q;
    MKAverage(int m, int k) {
        this->m = m;
        this->k = k;
        //Initialize both the cnt and s arrays to 0
        memset(cnt,0,sizeof cnt);
        memset(s,0,sizeof s);
    }

    //Returns the least significant 1 of x
    //For example x = 10111100
    //just return 100 which is 4
    int lowbit(int x){
        return x & -x;
    }
    //add an element
    void add(LL c[],int idx,int x){
        for(int i = idx;i <= N;i += lowbit(i)) c[i] += x;
    }
    //Find the prefix sum of 1 ~ idx
    LL preSum(LL c[],int idx){
        LL ans = 0;
        for(LL i = idx;i;i -= lowbit(i)) ans += c[i];
        return ans;
    }
    
    void addElement(int x) {
       //Insert element x, s records the sum of inserted elements, cnt records the number of inserted elements, that is, add 1 each time
        q.push(x);
        add(s,x,x);
        add(cnt,x,1);
        
        //If greater than m, pop an element from the head
        if(q.size() > this->m){
            int y = q.front();
            //queue head dequeue
            q.pop();
            //The following operation is equivalent to deleting this y in cnt and s
            add(s,y,-y);
            add(cnt,y,-1);
        }
    }

    LL getSum(int t){
        int l = 1,r = N;
        //Use the binary method to quickly find the first element whose prefix sum (here is the number of occurrences of the element) is greater than or equal to t
        //In other words, it is to find the sum of the first t elements (the elements are arranged from small to large)
        while(l < r){
            int mid = (l + r) >> 1;
            if(preSum(cnt,mid) >= t) r = mid;
            else l = mid + 1;
        }
        //At this time, r is the prefix of the first degree and the element greater than or equal to t
        //The latter item (preSum(cnt,r) - t) * 1LL * r, representing the redundant elements
        //For example, assume that the elements of [2,4,4,4,5] are inserted in sequence, and t = 2 at this time
        //So we're looking for the 2nd element, which is r = 4
        //But at this time preSum(s,4) is actually equal to 18, which is actually the prefix sum of four elements
        //Because 4 appears three times, we need to subtract the sum of the extra times
        //preSum(cnt,4) = 4
        //preSum(cnt,4) - t(t = 2) = 2
        //So the actual sum of the first t elements should be preSum(s,4) - (preSum(cnt,4) - t)*4 = 6
        //The reason for *1LL here is to convert it to long long to prevent overflow
        LL ans = preSum(s,r) - (preSum(cnt,r) - t) * 1LL * r;
        return ans;
    }
    
    int calculateMKAverage() {
        if(q.size() < m) return -1;
        else{
             //getSum(t) returns the sum of the first t numbers
            //Suppose at this time m = 6, k = 2 q = [2,3,3,5,5,9]
            //Observe that the answer should be the average of the sum of the middle two numbers
            //sum = getSum(4) - getSum(2) = 3 + 5 = 8;
            //So the answer ans = sum / (6 - 2 * 2) = 4
            LL ans = (getSum(m - k) - getSum(k)) / (m - 2 * k);
            return ans;
        }
    }
};

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