Skip to content
Related Articles
Open in App
Not now

Related Articles

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

Improve Article
Save Article
  • Last Updated : 20 Jan, 2023
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;
}


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<Integer> 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

Start Your Coding Journey Now!