Queries to calculate average of an array after removing K smallest and largest elements with updates

// C++ program for the above approach

#include <bits/stdc++.h>
using namespace std;

// Special Average class
class SpecialAverage {
public:
    // Three multisets to store the
    // smaller K larger K and the
    // remaining (n-2*k) integers
    multiset<int> left, mid, right;

    int n, k;

    // Stores the index of current
    // integer in running stream of
    // integers and the sum of integers
    // in mid multiset
    long pos, sum;

    // Array to store last n integers
    vector<int> v;

    // Constructor to initialize the
    // values of n and k and initialize
    // default values
    SpecialAverage(int nvalue, int kvalue)
    {
        n = nvalue;
        k = kvalue;
        pos = 0;
        sum = 0;
        for (int i = 0; i < n; i++)
            v.push_back(0);
    }

    // Add current integer in the
    // multiset left, mid or right
    // based on its value
    void add(int num)
    {
        // Firstly, insert num in the
        // left multiset
        left.insert(num);

        // Check if size of the left
        // multiset is greater than k
        if (left.size() > k) {

            // Stores the value of the
            // largest integer in the
            // left multiset
            int temp = *(prev(end(left)));

            // Insert temp in the
            // mid multiset
            mid.insert(temp);

            // Remove temp from the
            // left multiset
            left.erase(prev(end(left)));

            // Add the value of temp
            // to the current sum
            sum += temp;
        }

        // Check if the size of mid
        // multiset exceeds (n-2*k)
        if (mid.size() > (n - 2 * k)) {

            // Stores the value of the
            // largest integer in the
            // mid multiset
            int temp = *(prev(end(mid)));

            // Insert temp in the
            // right multiset
            right.insert(temp);

            // Remove temp from the
            // mid multiset
            mid.erase(prev(end(mid)));

            // Subtract the value of
            // temp from current sum
            sum -= temp;
        }
    }

    // Remove integer from the
    // multiset left, mid or right
    // based on its value
    void remove(int ele)
    {

        // If integer exists in the
        // left multiset
        if (ele <= *(left.rbegin()))
            left.erase(left.find(ele));

        // If integer exists in the
        // mid multiset
        else if (ele <= *(mid.rbegin())) {

            // Subtract the value of
            // integer from current sum
            sum -= ele;
            mid.erase(mid.find(ele));
        }

        // If integer exists in the
        // right multiset
        else
            right.erase(right.find(ele));

        // Balance all the multisets
        // left, right and mid check
        // if size of the left multiset
        // is less than k
        if (left.size() < k) {

            // Stores the value of
            // smallest integer
            // in mid multiset
            int temp = *(mid.begin());

            // Insert temp in the
            // left multiset
            left.insert(temp);

            // Remove temp from the
            // mid multiset
            mid.erase(mid.begin());

            // Subtract the value of
            // temp from current sum
            sum -= temp;
        }

        // Check if the size of mid
        // multiset becomes lesser
        // than (n-2*k)
        if (mid.size() < (n - 2 * k)) {

            // Stores the value of
            // smallest integer in
            // right multiset
            int temp = *(right.begin());

            // Insert temp in the
            // mid multiset
            mid.insert(temp);

            // Remove temp from the
            // right multiset
            right.erase(right.begin());

            // Add the value of temp
            // to the current sum
            sum += temp;
        }
    }

    // Function to add integer to
    // the current stream
    void addInteger(int num)
    {

        // Check if the total number
        // of elements in stream > n
        if (pos >= n)

            // Call the function to
            // remove  extra integer
            remove(v[pos % n]);

        // Insert the current integer
        // in the array v
        v[(pos++) % n] = num;

        // Call the function to add
        // the current integer in left,
        // mid or right multiset
        // based on its value
        add(num);
    }

    // Function to calculate the
    // special average
    int calculateSpecialAverage()
    {
        // Check if the total number
        // of elements is greater than
        // or equal to n
        if (pos >= n)

            // Return desired average
            return ((double)sum) / (n - 2 * k);
        return -1;
    }
};

// Function to process all the queries
void processQueries(int n, int k)
{
    // Create an object of the class
    SpecialAverage* avg
        = new SpecialAverage(n, k);

    // Add integer Query
    avg->addInteger(4);
    avg->addInteger(2);

    // Find average Query
    cout << avg->calculateSpecialAverage()
         << endl;

    // Add integer Query
    avg->addInteger(10);

    // Find average Query
    cout << avg->calculateSpecialAverage()
         << endl;
}

// Driver Code
int main()
{
    int N = 3, K = 1;
    processQueries(N, K);

    return 0;
}