Open In App

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

Improve
Improve
Like Article
Like
Save
Share
Report

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


Javascript




const mxN = 1e5;
  
let prime = new Array(mxN + 1).fill(true);
  
function sieveOfEratosthenes() {
    for (let p = 2; p * p <= mxN; p++) {
        if (prime[p]) {
            for (let i = p * p; i <= mxN; i += p) {
                prime[i] = false;
            }
        }
    }
}
  
function countSubarray(arr, n, k) {
    let s = new Set();
  
    let sum = 0;
    for (let i = 0; i < k; i++) {
        s.add(arr[i]);
        sum += arr[i];
    }
  
    let avgsum = Math.floor(sum / k);
  
    let ans = 0;
  
    let med = [...s].sort((a, b) => a - b)[(Math.floor((k + 1) / 2) - 1)];
  
    if (avgsum - med >= 0
        && ((prime[med] == false
             && prime[avgsum] == false)
            || (prime[med] != false
                && prime[avgsum] != false))) {
        ans++;
    }
  
    for (let i = 0; i < n - k; i++) {
        s.delete(arr[i]);
        s.add(arr[i + k]);
        sum -= arr[i];
        sum += arr[i + k];
        avgsum = Math.floor(sum / k);
        med = [...s].sort((a, b) => a - b)[(Math.floor((k + 1) / 2) - 1)];
  
        if (avgsum - med >= 0
            && ((prime[med] == false
                 && prime[avgsum] == false)
                || (prime[med] != false
                    && prime[avgsum] != false))) {
            ans++;
        }
    }
  
    return ans;
}
  
sieveOfEratosthenes();
  
let arr = [2, 4, 3, 5, 6];
let K = 3;
let N = arr.length;
  
console.log(countSubarray(arr, N, K));
// this code is contributed by devendra


Python3




import math
  
mxN = 100000
  
prime = [True] * (mxN + 1)
  
def sieveOfEratosthenes():
    for i in range(2, int(math.sqrt(mxN)) + 1):
        if prime[i]:
            for j in range(i * i, mxN + 1, i):
                prime[j] = False
  
def countSubarray(arr, n, k):
    s = set()
  
    sum_arr = sum(arr[:k])
    for i in range(k):
        s.add(arr[i])
  
    avgsum = sum_arr // k
  
    ans = 0
  
    med = sorted(s)[(k + 1) // 2 - 1]
  
    if avgsum - med >= 0 and ((not prime[med] and not prime[avgsum]) or (prime[med] and prime[avgsum])):
        ans += 1
  
    for i in range(n - k):
        s.remove(arr[i])
        s.add(arr[i + k])
        sum_arr = sum_arr - arr[i] + arr[i + k]
        avgsum = sum_arr // k
        med = sorted(s)[(k + 1) // 2 - 1]
  
        if avgsum - med >= 0 and ((not prime[med] and not prime[avgsum]) or (prime[med] and prime[avgsum])):
            ans += 1
  
    return ans
  
def main():
    sieveOfEratosthenes()
  
    arr = [2, 4, 3, 5, 6]
    K = 3
    N = len(arr)
  
    print(countSubarray(arr, N, K))
  
if __name__ == '__main__':
    main()


Output: 

2

 

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



Last Updated : 24 Mar, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads