# Count K-length subarrays whose average exceeds the median of the given array

• Last Updated : 25 Mar, 2022

Given an array arr[] consisting of N integers and a positive integer K, the task is to find the number of subarrays of size K whose average is greater than its median and both the average, median must be either prime or non-prime.

Examples:

Input: arr[] = {2, 4, 3, 5, 6}, K = 3
Output: 2
Explanation:
Following are the subarrays that satisfy the given conditions:

1. {2, 4, 3}: The median of this subarray is 3, and the average is (2 + 4 + 3)/3 = 3. As, both the median and average are prime and average >= median. So the count this subarray.
2. {4, 3, 5}: The median of this subarray is 4, and the average is (4 + 3 + 5)/3 = 4. As, both the median and average are non-prime and average >= median. So the count this subarray.

Therefore, the total number of subarrays are 2.

Input: arr[] = {2, 4, 3, 5, 6}, K = 2
Output: 3

Approach: The given problem can be solved using Policy-based Data Structures i.e., ordered_set. Follow the steps below to solve the given problem:

• Precompute all the primes and non-primes till 105 using Sieve Of Eratosthenes.
• Initialize a variable, say count that stores the resultant count of subarrays.
• Find the average and median of the first K elements and if the average >= median and both average and medians are either prime or non-prime, then increment the count by 1.
• Store the first K array elements in the ordered_set.
• Traverse the given array over the range [0, N – K] and perform the following steps:
• Remove the current element arr[i] from the ordered_set and add (i + k)th element i.e., arr[i + K] to the ordered_set.
• Find the median of the array using the function find_order_by_set((K + 1)/2 – 1).
• Find the average of the current subarray.
• If the average >= median and both average and medians are either prime or non-prime, then increment the count by 1.
• After completing the above steps, print the value of count as the result.

Below is the implementation of the above approach.

## C++

 `// C++ program for the above approach` `#include ``#include ``#include ``#include ``using` `namespace` `__gnu_pbds;` `using` `namespace` `std;``typedef` `tree<``int``, null_type, less_equal<``int``>,``             ``rb_tree_tag,``             ``tree_order_statistics_node_update>``    ``ordered_set;` `const` `int` `mxN = (``int``)1e5;` `// Stores whether i is prime or not``bool` `prime[mxN + 1];` `// Function to precompute all the prime``// numbers using sieve of eratosthenes``void` `SieveOfEratosthenes()``{``    ``// Initialize the prime array``    ``memset``(prime, ``true``, ``sizeof``(prime));` `    ``// Iterate over the range [2, mxN]``    ``for` `(``int` `p = 2; p * p <= mxN; p++) {` `        ``// If the prime[p] is unchanged,``        ``// then it is a prime``        ``if` `(prime[p]) {` `            ``// Mark all multiples of p``            ``// as non-prime``            ``for` `(``int` `i = p * p;``                 ``i <= mxN; i += p)``                ``prime[i] = ``false``;``        ``}``    ``}``}` `// Function to find number of subarrays``// that satisfy the given criteria``int` `countSubarray(``int` `arr[], ``int` `n, ``int` `k)``{``    ``// Initialize the ordered_set``    ``ordered_set s;` `    ``// Stores the sum for subarray``    ``int` `sum = 0;``    ``for` `(``int` `i = 0; i < (``int``)k; i++) {``        ``s.insert(arr[i]);``        ``sum += arr[i];``    ``}` `    ``// Stores the average for each``    ``// possible subarray``    ``int` `avgsum = sum / k;` `    ``// Stores the count of subarrays``    ``int` `ans = 0;` `    ``// For finding the median use the``    ``// find_by_order(k) that returns``    ``// an iterator to kth element``    ``int` `med = *s.find_by_order(``        ``(k + 1) / 2 - 1);` `    ``// Check for the valid condition``    ``if` `(avgsum - med >= 0``        ``&& ((prime[med] == 0``             ``&& prime[avgsum] == 0)``            ``|| (prime[med] != 0``                ``&& prime[avgsum] != 0))) {` `        ``// Increment the resultant``        ``// count of subarray``        ``ans++;``    ``}` `    ``// Iterate the subarray over the``    ``// the range [0, N - K]``    ``for` `(``int` `i = 0; i < (``int``)(n - k); i++) {` `        ``// Erase the current element``        ``// arr[i]``        ``s.erase(s.find_by_order(``            ``s.order_of_key(arr[i])));` `        ``// The function Order_of_key(k)``        ``// returns the number of items``        ``// that are strictly smaller``        ``// than K``        ``s.insert(arr[i + k]);``        ``sum -= arr[i];` `        ``// Add the (i + k)th element``        ``sum += arr[i + k];` `        ``// Find the average``        ``avgsum = sum / k;` `        ``// Get the median value``        ``med = *s.find_by_order(``            ``(k + 1) / 2 - 1);` `        ``// Check the condition``        ``if` `(avgsum - med >= 0``            ``&& ((prime[med] == 0``                 ``&& prime[avgsum] == 0)``                ``|| (prime[med] != 0``                    ``&& prime[avgsum] != 0))) {` `            ``// Increment the count of``            ``// subarray``            ``ans++;``        ``}``    ``}` `    ``// Return the resultant count``    ``// of subarrays``    ``return` `ans;``}` `// Driver Code``int` `main()``{``    ``// Precompute all the primes``    ``SieveOfEratosthenes();` `    ``int` `arr[] = { 2, 4, 3, 5, 6 };``    ``int` `K = 3;``    ``int` `N = ``sizeof``(arr) / ``sizeof``(arr[0]);` `    ``cout << countSubarray(arr, N, K);` `    ``return` `0;``}`

Output:

`2`

Time Complexity: O(N*log N + N*log(log N))
Auxiliary Space: O(N)

My Personal Notes arrow_drop_up