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

• Last Updated : 20 Jan, 2023

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

## Java

 `import` `java.util.*;` `import` `java.util.stream.*;`   `public` `class` `Gfg {` `    ``static` `final` `int` `mxN = (``int``)1e5;`   `    ``static` `boolean``[] prime = ``new` `boolean``[mxN + ``1``];`   `    ``static` `void` `sieveOfEratosthenes()` `    ``{` `        ``Arrays.fill(prime, ``true``);`   `        ``for` `(``int` `p = ``2``; p * p <= mxN; p++) {` `            ``if` `(prime[p]) {` `                ``for` `(``int` `i = p * p; i <= mxN; i += p) {` `                    ``prime[i] = ``false``;` `                ``}` `            ``}` `        ``}` `    ``}`   `    ``static` `int` `countSubarray(``int``[] arr, ``int` `n, ``int` `k)` `    ``{` `        ``TreeSet s = ``new` `TreeSet<>();`   `        ``int` `sum = ``0``;` `        ``for` `(``int` `i = ``0``; i < k; i++) {` `            ``s.add(arr[i]);` `            ``sum += arr[i];` `        ``}`   `        ``int` `avgsum = sum / k;`   `        ``int` `ans = ``0``;`   `        ``int` `med = s.stream()` `                      ``.skip((k + ``1``) / ``2` `- ``1``)` `                      ``.findFirst()` `                      ``.get();`   `        ``if` `(avgsum - med >= ``0` `            ``&& ((prime[med] == ``false` `                 ``&& prime[avgsum] == ``false``)` `                ``|| (prime[med] != ``false` `                    ``&& prime[avgsum] != ``false``))) {` `            ``ans++;` `        ``}`   `        ``for` `(``int` `i = ``0``; i < n - k; i++) {` `            ``s.remove(arr[i]);` `            ``s.add(arr[i + k]);` `            ``sum -= arr[i];` `            ``sum += arr[i + k];` `            ``avgsum = sum / k;` `            ``med = s.stream()` `                      ``.skip((k + ``1``) / ``2` `- ``1``)` `                      ``.findFirst()` `                      ``.get();`   `            ``if` `(avgsum - med >= ``0` `                ``&& ((prime[med] == ``false` `                     ``&& prime[avgsum] == ``false``)` `                    ``|| (prime[med] != ``false` `                        ``&& prime[avgsum] != ``false``))) {` `                ``ans++;` `            ``}` `        ``}`   `        ``return` `ans;` `    ``}`   `    ``public` `static` `void` `main(String[] args)` `    ``{` `        ``sieveOfEratosthenes();`   `        ``int``[] arr = { ``2``, ``4``, ``3``, ``5``, ``6` `};` `        ``int` `K = ``3``;` `        ``int` `N = arr.length;`   `        ``System.out.println(countSubarray(arr, N, K));` `    ``}` `}`

## C#

 `using` `System;` `using` `System.Linq;` `using` `System.Collections.Generic;`   `class` `Gfg {` `    ``static` `int` `mxN = 100000;`   `    ``static` `bool``[] prime = ``new` `bool``[mxN + 1];`   `    ``static` `void` `sieveOfEratosthenes()` `    ``{` `        ``for` `(``int` `i = 0; i < prime.Length; i++) {` `            ``prime[i] = ``true``;` `        ``}`   `        ``for` `(``int` `p = 2; p * p <= mxN; p++) {` `            ``if` `(prime[p]) {` `                ``for` `(``int` `i = p * p; i <= mxN; i += p) {` `                    ``prime[i] = ``false``;` `                ``}` `            ``}` `        ``}` `    ``}`   `    ``static` `int` `countSubarray(``int``[] arr, ``int` `n, ``int` `k)` `    ``{` `        ``SortedSet<``int``> s = ``new` `SortedSet<``int``>();`   `        ``int` `sum = 0;` `        ``for` `(``int` `i = 0; i < k; i++) {` `            ``s.Add(arr[i]);` `            ``sum += arr[i];` `        ``}`   `        ``int` `avgsum = sum / k;`   `        ``int` `ans = 0;`   `        ``int` `med = s.Skip((k + 1) / 2 - 1).First();`   `        ``if` `(avgsum - med >= 0` `            ``&& ((prime[med] == ``false` `                 ``&& prime[avgsum] == ``false``)` `                ``|| (prime[med] != ``false` `                    ``&& prime[avgsum] != ``false``))) {` `            ``ans++;` `        ``}`   `        ``for` `(``int` `i = 0; i < n - k; i++) {` `            ``s.Remove(arr[i]);` `            ``s.Add(arr[i + k]);` `            ``sum -= arr[i];` `            ``sum += arr[i + k];` `            ``avgsum = sum / k;` `            ``med = s.Skip((k + 1) / 2 - 1).First();`   `            ``if` `(avgsum - med >= 0` `                ``&& ((prime[med] == ``false` `                     ``&& prime[avgsum] == ``false``)` `                    ``|| (prime[med] != ``false` `                        ``&& prime[avgsum] != ``false``))) {` `                ``ans++;` `            ``}` `        ``}`   `        ``return` `ans;` `    ``}`   `    ``public` `static` `void` `Main(``string``[] args)` `    ``{` `        ``sieveOfEratosthenes();`   `        ``int``[] arr = { 2, 4, 3, 5, 6 };` `        ``int` `K = 3;` `        ``int` `N = arr.Length;`   `        ``Console.WriteLine(countSubarray(arr, N, K));` `    ``}` `}`

Output:

`2`

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

My Personal Notes arrow_drop_up
Related Articles