Related Articles

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

• Last Updated : 27 Jul, 2021

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:

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with experts, please refer DSA Live Classes for Working Professionals and Competitive Programming Live for Students.

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];`` ` `// Fuction to precompute all the prime``// numbers using sieve of erastothenes``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);`` ` `    ``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