Average of max K numbers in a stream

Given a list of ‘N’ numbers, and an integer ‘K’. The task is to print the average of max ‘K’ numbers after each query where a query consists of an integer element which needs to be added to the list of elements.
Note: The queries are defined with an integer array ‘q’

Examples:

Input: N = 4, K = 3, arr = {1, 2, 3, 4}, q = {7, 2, 1, 5}
Output: 4.666666
4.666666
4.666666
5.333333
After query 1, arr = {1, 2, 3, 4, 7} and the average of max K (i.e. {3, 4, 7}) elements is 4.666666.
After query 2, arr = {1, 2, 3, 4, 7, 2} and the average is 4.666666 for {3, 4, 7}.
After query 3, arr = {1, 2, 3, 4, 7, 2, 1} and the average is 4.666666 for {3, 4, 7}.
After query 4, arr = {1, 2, 3, 4, 7, 2, 5} and the average is 5.333333 for {4, 5, 7}.

Input: N = 5, K = 4, arr = {1, 2, 2, 3, 3}, q = {2, 5, 1}
Output: 2.5
3.25
3.25



Approach: Heap (Min Heap) data structure can be used to solve problems like these where insertion and deletions of the elements can be performed in O(log n) time.

  • Initially, store the max k elements from the given list of elements in the min heap.
  • If the incoming element is less than or equal to the element currently at the root of the min heap then discard the element as it’ll have no effect on the average.
  • Else if, the number is greater than the root element then remove the root of the min heap followed by an insertion of the new element and then calculate the average of the elements currently in the heap.
  • Print the average and repeat the above two steps for all incoming elements.

Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the approach
import java.util.*;
  
class GFG {
  
    // Function that returns the
    // average of max k elements in
    // the list after each query
    static void max_average_k_numbers(int n,
                                      int k,
                                      int m,
                                      int[] arr,
                                      int[] query)
    {
        double max_avg = 0.0;
  
        // min-heap to maintain
        // the max k elements at
        // any point of time;
        PriorityQueue<Integer> pq = new PriorityQueue<Integer>();
  
        // Sort the array
        // in ascending order
        Arrays.sort(arr);
  
        // add max k elements
        // to the heap
        double sum = 0;
        for (int i = n - 1; i >= n - k; i--) {
            pq.add(arr[i]);
            sum = sum + arr[i];
        }
  
        // perform offline queries
        for (int i = 0; i < m; i++) {
  
            // if the minimum element in
            // the heap is less than
            // the incoming element
            if (query[i] > pq.peek()) {
                int polled = pq.poll();
                pq.add(query[i]);
  
                // decrement the current
                // sum by the polled element
                sum = sum - polled;
  
                // increment sum by the
                // incoming element
                sum = sum + query[i];
            }
  
            // compute the average
            max_avg = sum / (double)k;
            System.out.println(max_avg);
        }
    }
  
    // Driver code
    public static void main(String[] args)
    {
        int n = 4;
        int k = 3;
        int m = 4;
        int[] arr = new int[] { 1, 2, 3, 4 };
        int[] query = new int[] { 7, 2, 1, 5 };
  
        max_average_k_numbers(n, k, m, arr, query);
    }
}

chevron_right


Output:

4.666666666666667
4.666666666666667
4.666666666666667
5.333333333333333


My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.