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:
- {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.
- {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++ 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;
} |
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));
}
} |
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));
}
} |
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 |
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()
|
2
Time Complexity: O(N*log N + N*log(log N))
Auxiliary Space: O(N)