Skip to content
Related Articles

Related Articles

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

View Discussion
Improve Article
Save Article
  • Last Updated : 25 Mar, 2022
View Discussion
Improve Article
Save Article

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 <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <stdlib.h>
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
Recommended Articles
Page :

Start Your Coding Journey Now!