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

Given two positive integers N and K, initialize an empty array arr[] and Q number of queries of the following two types:

addInteger(x): Insert element X in the array arr[]. If the size of the array becomes greater than N, then remove the element from the beginning of the array.
calculateSpecialAverage(): Find the average of array elements after removing the first K smallest and the largest elements. If the size of the array is less than N, then print -1.

The task is to process the given queries and print the results accordingly.

Examples:

Input: N = 3, K = 1, Q = 5, Queries[] = { addInteger(4), addInteger(2), calculateSpecialAverage(), addInteger(10), calculateSpecialAverage() }
Output: -1, 4
Explanation:
Below are the queries performed:
Query 1 is to insert element 4 in the array, arr[] becomes {4}.
Query 2 is to insert the element 2 in the array, arr[] becomes {4, 2}.
Query 3 is to find the average. Since the size of the array is less than N(= 3), print -1.
Query 4 is to insert the element 10 in the array, arr[] becomes {4, 2, 10}.
Query 5 is to find the average. Since the size of the array =3, remove K(= 1) the smallest element i.e., 2 and then remove the largest element i.e., 10. Now the average will be 4/1 = 4.

Input: N = 3, K = 1, Q = 5, Queries[] = { addInteger(4), addInteger(21), calculateSpecialAverage() }
Output: -1

Approach: The given problem can be solved by using the multiset. Follow the steps below to solve the problem:

Initialize three multisets left, right and mid to store K smallest, K largest, and the remaining (N – 2*K) integers.
Declare a vector v of size N to store the last N integers.
Declare an integer variable pos to keep track of the current index and an integer sum to store the sum of values in the mid multiset, which is the desired sum for calculating the average of (N – 2*K) integers.
Define a function add to insert an integer in left, mid, or right multiset based on its value. To insert integer in the correct multiset follow the following steps:
- Initially, insert the integer in the left multiset.
- If the size of the left multiset exceeds k, delete the largest integer from the left multiset and insert it in the mid multiset.
- If the size of mid multiset exceeds (n – 2*k), delete the largest integer from mid multiset and insert it in the right multiset.
- While removing and adding integers in mid multiset maintain the correct value of sum i.e. if an integer is removed from mid multiset subtract its value from the sum and similarly if an integer is added in the mid multiset, add its value to sum, to keep the value of sum updated.
Declare a function remove to erase the integers from left, right, or mid multiset.
- If the value of the integer num is less than or equal to the largest integer in the left multiset, then find and erase the integer from the left multiset, otherwise search for it in mid or right multiset, based on its value.
- If after the removal of num, the size of the left multiset decreases, delete the smallest integer from the mid multiset and insert it in the left multiset.
- Similarly, for mid multiset, if the size decreases, delete the smallest integer from the right multiset and insert it in the mid multiset.
Define a function addInteger to add an integer to the current stream:
- If the number of integers in the current stream(sum of the size of left, mid, right multiset) exceeds n, remove the integer at index pos%n.
- Call the function add to insert the integer in the left, mid, or right multiset, based on its value.
Define the function calculateSpecialAverage to return the special average:
- If the number of integers in the stream is less than N, then print -1.
- Otherwise, print the average as (sum/(N – 2*K)).

Below is the implementation of the above approach:

C++

// 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;
}

Java

// Java program for the above approach

import java.util.*;
 
class SpecialAverage {

    // Three TreeSets to store the

    // smaller K larger K and the

    // remaining (n-2*k) integers

    TreeSet<Integer> left, mid, right;
 
    int n, k;
 
    // Stores the index of current

    // integer in running stream of

    // integers and the sum of integers

    // in mid TreeSet

    long pos, sum;
 
    // Array to store last n integers

    ArrayList<Integer> 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;

        v = new ArrayList<>();

        for (int i = 0; i < n; i++)

            v.add(0);

        left = new TreeSet<>();

        mid = new TreeSet<>();

        right = new TreeSet<>();

    }
 
    // Add current integer in the

    // TreeSet left, mid or right

    // based on its value

    void add(int num) {

        // Firstly, insert num in the

        // left TreeSet

        left.add(num);
 
        // Check if size of the left

        // TreeSet is greater than k

        if (left.size() > k) {

            // Stores the value of the

            // largest integer in the

            // left TreeSet

            int temp = left.pollLast();
 
            // Insert temp in the

            // mid TreeSet

            mid.add(temp);
 
            // Add the value of temp

            // to the current sum

            sum += temp;

        }
 
        // Check if the size of mid

        // TreeSet exceeds (n-2*k)

        if (mid.size() > (n - 2 * k)) {

            // Stores the value of the

            // largest integer in the

            // mid TreeSet

            int temp = mid.pollLast();
 
            // Insert temp in the

            // right TreeSet

            right.add(temp);
 
            // Subtract the value of

            // temp from current sum

            sum -= temp;

        }

    }
 
    // Remove integer from the

    // TreeSet left, mid or right

    // based on its value

    void remove(int ele) {

        // If integer exists in the

        // left TreeSet

        if (ele <= left.last())

            left.remove(ele);
 
        // If integer exists in the

        // mid TreeSet

        else if (ele <= mid.last()) {

            // Subtract the value of

            // integer from current sum

            sum -= ele;

            mid.remove(ele);

        }
 
        // If integer exists in the

        // right TreeSet

        else

            right.remove(ele);
 
        // Balance all the TreeSets

        // left, right and mid check

        // if size of the left TreeSet

        // is less than k

        if (left.size() < k) {

            // Stores the value of

            // smallest integer

            // in mid TreeSet

            int temp = mid.first();
 
            // Insert temp in the

            // left TreeSet

            left.add(temp);
 
            // Subtract the value of

            // temp from current sum

            sum -= temp;
 
            // Remove temp from the

            // mid TreeSet

            mid.remove(temp);

        }
 
        // Check if the size of mid

        // TreeSet becomes lesser

        // than (n-2*k)

        if (mid.size() < (n - 2 * k)) {

            // Stores the value of

            // smallest integer in

            // right TreeSet

            int temp = right.first();
 
            // Insert temp in the

            // mid TreeSet

            mid.add(temp);
 
            // Add the value of temp

            // to the current sum

            sum += temp;
 
            // Remove temp from the

            // right TreeSet

            right.remove(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.get((int) (pos % n)));
 
        // Insert the current integer

        // in the array v

        v.set((int) (pos++ % n), num);
 
        // Call the function to add

        // the current integer in left,

        // mid or right TreeSet

        // 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 ((int) sum) / (n - 2 * k);

        return -1;

    }

    // Function to process all the queries

    static 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

        System.out.println(avg.calculateSpecialAverage());

        // Add integer Query

        avg.addInteger(10);

        // Find average Query

        System.out.println(avg.calculateSpecialAverage());

    }

    // Driver Code

    public static void main(String[] args) {

        int N = 3, K = 1;

        processQueries(N, K);

    }
}

Python3

class SpecialAverage:

    def __init__(self, nvalue, kvalue):

        # Three sets to store the smaller K, larger K, and the remaining (n-2*k) integers

        self.left = set()

        self.mid = set()

        self.right = set()

        self.n = nvalue

        self.k = kvalue

        self.pos = 0

        self.sum = 0

        self.v = [0] * nvalue
 
    # Add the current integer to the SortedSets left, mid, or right based on its value

    def add(self, num):

        # Firstly, insert num into the left SortedSet

        self.left.add(num)

        # Check if the size of the left SortedSet is greater than k

        if len(self.left) > self.k:

            # Stores the value of the largest integer in the left SortedSet
 
            temp = max(self.left)

            # Insert temp into the mid SortedSet
 
            self.mid.add(temp)

            # Add the value of temp to the current sum
 
            self.sum += temp

            # Remove temp from the left SortedSet
 
            self.left.remove(temp)

        # Check if the size of the mid SortedSet exceeds (n-2*k)
 
        if len(self.mid) > (self.n - 2 * self.k):

            # Stores the value of the largest integer in the mid SortedSet
 
            temp = max(self.mid)

            # Insert temp into the right SortedSet
 
            self.right.add(temp)

            # Subtract the value of temp from the current sum
 
            self.sum -= temp
 
            # Remove temp from the mid SortedSet
 
            self.mid.remove(temp)

    # Remove the integer from the SortedSets left, mid, or right based on its value
 
    def remove(self, ele):

        # If the integer exists in the left SortedSet

        if ele <= max(self.left, default=float('-inf')):

            self.left.remove(ele)

        # If the integer exists in the mid SortedSet

        elif ele <= max(self.mid, default=float('-inf')):

            # Subtract the value of the integer from the current sum

            self.sum -= ele

            self.mid.remove(ele)

        # If the integer exists in the right SortedSet

        else:

            self.right.remove(ele)
 
        # Balance all the SortedSets left, right, and mid

        # Check if the size of the left SortedSet is less than k

        if len(self.left) < self.k:

            # Stores the value of the smallest integer in the mid SortedSet

            temp = min(self.mid, default=float('inf'))

            # Insert temp into the left SortedSet

            self.left.add(temp)

            # Subtract the value of temp from the current sum

            self.sum -= temp

            # Remove temp from the mid SortedSet

            self.mid.remove(temp)
 
        # Check if the size of the mid SortedSet becomes less than (n-2*k)

        if len(self.mid) < (self.n - 2 * self.k):

            # Stores the value of the smallest integer in the right SortedSet

            temp = min(self.right, default=float('inf'))

            # Insert temp into the mid SortedSet

            self.mid.add(temp)

            # Add the value of temp to the current sum

            self.sum += temp

            # Remove temp from the right SortedSet

            self.right.remove(temp)
 
    # Function to add an integer to the current stream

    def add_integer(self, num):

        # Check if the total number of elements in the stream is greater than n

        if self.pos >= self.n:

            # Call the function to remove the extra integer

            self.remove(self.v[int(self.pos % self.n)])
 
        # Insert the current integer into the array v

        self.v[int(self.pos % self.n)] = num
 
        # Call the function to add the current integer into left, mid, or right SortedSets

        # based on its value

        self.add(num)

        self.pos += 1

    # Function to calculate the special average
 
    def calculate_special_average(self):

        # Check if the total number of elements is greater than or equal to n

        if self.pos >= self.n:

            # Return the desired average

            return int(self.sum / (self.n - 2 * self.k))

        return -1
 
# Function to process all the queries
 
def process_queries(n, k):

    # Create an object of the SpecialAverage class

    avg = SpecialAverage(n, k)
 
    # Add integer Query

    avg.add_integer(4)

    avg.add_integer(2)
 
    # Find average Query

    print(avg.calculate_special_average())
 
    # Add integer Query

    avg.add_integer(10)
 
    # Find average Query

    print(avg.calculate_special_average())
 
if __name__ == "__main__":

    N, K = 3, 1

    process_queries(N, K)

using System;

using System.Collections.Generic;
 
class SpecialAverage
{

    // Three SortedSets to store the smaller K, larger K, and the remaining (n-2*k) integers

    SortedSet<int> left, mid, right;
 
    int n, k;
 
    // Stores the index of the current integer in the running stream of integers

    // and the sum of integers in the mid SortedSet

    long pos, sum;
 
    // Array to store the last n integers

    List<int> v;
 
    // Constructor to initialize the values of n and k and initialize default values

    public SpecialAverage(int nvalue, int kvalue)

    {

        n = nvalue;

        k = kvalue;

        pos = 0;

        sum = 0;

        v = new List<int>();

        for (int i = 0; i < n; i++)

            v.Add(0);

        left = new SortedSet<int>();

        mid = new SortedSet<int>();

        right = new SortedSet<int>();

    }
 
    // Add the current integer to the SortedSets left, mid, or right based on its value

    public void Add(int num)

    {

        // Firstly, insert num into the left SortedSet

        left.Add(num);
 
        // Check if the size of the left SortedSet is greater than k

        if (left.Count > k)

        {

            // Stores the value of the largest integer in the left SortedSet

            int temp = left.Max;
 
            // Insert temp into the mid SortedSet

            mid.Add(temp);
 
            // Add the value of temp to the current sum

            sum += temp;
 
            // Remove temp from the left SortedSet

            left.Remove(temp);

        }
 
        // Check if the size of the mid SortedSet exceeds (n-2*k)

        if (mid.Count > (n - 2 * k))

        {

            // Stores the value of the largest integer in the mid SortedSet

            int temp = mid.Max;
 
            // Insert temp into the right SortedSet

            right.Add(temp);
 
            // Subtract the value of temp from the current sum

            sum -= temp;
 
            // Remove temp from the mid SortedSet

            mid.Remove(temp);

        }

    }
 
    // Remove the integer from the SortedSets left, mid, or right based on its value

    public void Remove(int ele)

    {

        // If the integer exists in the left SortedSet

        if (ele <= left.Max)

            left.Remove(ele);
 
        // If the integer exists in the mid SortedSet

        else if (ele <= mid.Max)

        {

            // Subtract the value of the integer from the current sum

            sum -= ele;

            mid.Remove(ele);

        }
 
        // If the integer exists in the right SortedSet

        else

            right.Remove(ele);
 
        // Balance all the SortedSets left, right, and mid

        // Check if the size of the left SortedSet is less than k

        if (left.Count < k)

        {

            // Stores the value of the smallest integer in the mid SortedSet

            int temp = mid.Min;
 
            // Insert temp into the left SortedSet

            left.Add(temp);
 
            // Subtract the value of temp from the current sum

            sum -= temp;
 
            // Remove temp from the mid SortedSet

            mid.Remove(temp);

        }
 
        // Check if the size of the mid SortedSet becomes less than (n-2*k)

        if (mid.Count < (n - 2 * k))

        {

            // Stores the value of the smallest integer in the right SortedSet

            int temp = right.Min;
 
            // Insert temp into the mid SortedSet

            mid.Add(temp);
 
            // Add the value of temp to the current sum

            sum += temp;
 
            // Remove temp from the right SortedSet

            right.Remove(temp);

        }

    }
 
    // Function to add an integer to the current stream

    public void AddInteger(int num)

    {

        // Check if the total number of elements in the stream is greater than n

        if (pos >= n)

            // Call the function to remove the extra integer

            Remove(v[(int)(pos % n)]);
 
        // Insert the current integer into the array v

        v[(int)(pos++ % n)] = num;
 
        // Call the function to add the current integer into left, mid, or right SortedSets

        // based on its value

        Add(num);

    }
 
    // Function to calculate the special average

    public int CalculateSpecialAverage()

    {

        // Check if the total number of elements is greater than or equal to n

        if (pos >= n)

            // Return the desired average

            return (int)sum / (n - 2 * k);

        return -1;

    }
 
    // Function to process all the queries

    public static void ProcessQueries(int n, int k)

    {

        // Create an object of the SpecialAverage class

        SpecialAverage avg = new SpecialAverage(n, k);
 
        // Add integer Query

        avg.AddInteger(4);

        avg.AddInteger(2);
 
        // Find average Query

        Console.WriteLine(avg.CalculateSpecialAverage());
 
        // Add integer Query

        avg.AddInteger(10);
 
        // Find average Query

        Console.WriteLine(avg.CalculateSpecialAverage());

    }
 
    // Driver Code

    public static void Main(string[] args)

    {

        int N = 3, K = 1;

        ProcessQueries(N, K);

    }
}

Javascript

class SpecialAverage {

  constructor(nvalue, kvalue) {

    // Three sets to store smaller K, larger K, and remaining integers

    this.left = new Set();

    this.mid = new Set();

    this.right = new Set();

    this.n = nvalue; // Total number of integers to consider

    this.k = kvalue; // Parameter value K

    this.pos = 0; // Position tracker

    this.sum = 0; // Sum of selected integers

    this.v = new Array(nvalue).fill(0); // Array to store the integers

  }
 
  // Add the current integer to the appropriate set

  add(num) {

    this.left.add(num);
 
    // Check and move integers from left to mid set if left set size exceeds K

    if (this.left.size > this.k) {

      let temp = Math.max(...this.left);

      this.mid.add(temp);

      this.sum += temp;

      this.left.delete(temp);

    }
 
    // Check and move integers from mid to right set if mid set size exceeds (n - 2 * K)

    if (this.mid.size > this.n - 2 * this.k) {

      let temp = Math.max(...this.mid);

      this.right.add(temp);

      this.sum -= temp;

      this.mid.delete(temp);

    }

  }
 
  // Remove the specified integer from the sets

  remove(ele) {

    if (ele <= Math.max(...this.left, Number.NEGATIVE_INFINITY)) {

      this.left.delete(ele);

    } else if (ele <= Math.max(...this.mid, Number.NEGATIVE_INFINITY)) {

      this.sum -= ele;

      this.mid.delete(ele);

    } else {

      this.right.delete(ele);

    }
 
    // Ensure left set size is maintained as K

    if (this.left.size < this.k) {

      let temp = Math.min(...this.mid, Number.POSITIVE_INFINITY);

      this.left.add(temp);

      this.sum -= temp;

      this.mid.delete(temp);

    }
 
    // Ensure mid set size is maintained as (n - 2 * K)

    if (this.mid.size < this.n - 2 * this.k) {

      let temp = Math.min(...this.right, Number.POSITIVE_INFINITY);

      this.mid.add(temp);

      this.sum += temp;

      this.right.delete(temp);

    }

  }
 
  // Add an integer to the stream

  addInteger(num) {

    // If total elements exceed n, remove the oldest element

    if (this.pos >= this.n) {

      this.remove(this.v[this.pos % this.n]);

    }
 
    // Add the current integer to the array and sets

    this.v[this.pos % this.n] = num;

    this.add(num);

    this.pos++;

  }
 
  // Calculate the special average

  calculateSpecialAverage() {

    // If total elements are sufficient, return the special average

    if (this.pos >= this.n) {

      return Math.floor(this.sum / (this.n - 2 * this.k));

    }

    return -1; // Insufficient elements to calculate the special average

  }
}
 
// Function to process queries

function processQueries(n, k) {

  let avg = new SpecialAverage(n, k);
 
  // Add integers to the stream

  avg.addInteger(4);

  avg.addInteger(2);

  console.log(avg.calculateSpecialAverage()); // Calculate and print the special average
 
  // Add another integer to the stream

  avg.addInteger(10);

  console.log(avg.calculateSpecialAverage()); // Calculate and print the updated special average
}
 
// Example usage:
const N = 3;
const K = 1;

processQueries(N, K); // Call the function to process queries