Author : tejasdeore7
Tester : adarsh_sinhg
Editorialist : tejasdeore7
DIFFICULTY :
Medium
PREREQUISITES :
Policy based data structures
PROBLEM :
You are given an array A[1……N] of length N.Your task is to find the number of K length subarrays such that average of the subarray is greater than or equal to the median value of that subarray in sorted format. Also the median and average must both be prime or non prime.
Median of an even length subarray would be the value at the smaller index of the two. For eg: Median of {2,3,4,5} would be 3, not 4. Consider the integer(floored) value of average for calculations. You have to answer T independent test cases.
EXPLANATION :
First calculate/ precompute all the primes and non-primes using SieveOfEratosthenes. After that, there are various possible solutions.
Two of them are mentioned below :
Solution 1 :
The problem can be solved using policy based data structures/ ordered_set in the following way :
Calculate the sum, average and median for the initial K elements.
Store these elements in an ordered_set.
Then loop for elements from i+k to n-1.
So first remove the i th element from the set and add (i+k) th element to the set.
The ordered_set will perform all operations in O(logn) time and hence the median value can easily be calculated from the set.
Then check for the prime/non-prime and average >= median conditions, if they are satisfied, add to the count, else don’t add to the count.
Solution 2 :
The problem can be solved using two maps in the following way :
Store the first half of the subarray in map1 and second half in map2.
Calculate the sum, average and median for the initial K elements.
Distribute these elements evenly in the two maps.
Then loop for elements from i+k to n-1.
So first remove the i th element accordingly (explained in Solution 2 code) from either map1 or map2.
Then add (i+k) th element accordingly (explained in Solution 2 code) in either map1 or map2.
Make changes in the subarray sum and average accordingly.
So, the median value for a subarray will always be the last value of map1.
Then check for the prime/non-prime and average >= median conditions, if they are satisfied, add to the count, else don’t add to the count.
SOLUTION CODES :
Solution 1/ Setter's Solution
#include <bits/stdc++.h>
// Header files, namespaces
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
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;
// Use less_equal<int> because numbers are not necessarily distinct
#define int long long
const int mxN = (int)1e5;
bool prime[mxN+1];
// Create a boolean array "prime[0..n]" and initialize
// all entries it as true. A value in prime[i] will
// finally be false if i is Not a prime, else true.
void SieveOfEratosthenes()
{
memset(prime, true, sizeof(prime));
for (int p=2; p*p<=mxN; p++)
{
// If prime[p] is not changed, then it is a prime
if (prime[p])
{
// Update all multiples of p
for (int i=p*p; i<=mxN; i += p)
prime[i] = false;
}
}
}
int32_t main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
SieveOfEratosthenes();
// Use Sieve to find(or precompute) all prime numbers upto 10^5
int tc;
cin >> tc;
while(tc--){
int n, k;
cin >> n >> k;
vector<int> arr(n);
for(int i = 0;i < (int)n; i++){
cin >> arr[i];
}
ordered_set s;
int sum = 0;// To store sum for every subarray
for(int i = 0;i < (int)k; i++) {
s.insert(arr[i]);
sum += arr[i];
}
int avgsum = sum/k;// To store average for every subarray
int ans = 0;// To store final count of subarrays
int med = *s.find_by_order((k+1)/2 - 1);// To store the median value
//find_by_order(k):
//It returns to an iterator to the kth element (counting from zero) in the set in O(logn) time.
if(avgsum-med>=0 && ((prime[med]==0 && prime[avgsum]==0) || (prime[med]!=0 && prime[avgsum]!=0))){
ans++;
}// Check the condition
for(int i = 0;i < (int)(n-k); i++){
s.erase(s.find_by_order(s.order_of_key(arr[i])));// erase arr[i]
//order_of_key(k) :
//It returns to the number of items that are strictly smaller than our item k in O(logn) time.
s.insert(arr[i+k]);// add arr[i+k]
sum -= arr[i];
sum += arr[i+k];
avgsum = sum/k;// Change the sum accordingly and calculate average
med = *s.find_by_order((k+1)/2 - 1);// get the median value
if(avgsum-med>=0 && ((prime[med]==0 && prime[avgsum]==0) || (prime[med]!=0 && prime[avgsum]!=0))){
ans++;
}// Check the condition
}
cout << ans << "\n";
}
return 0;
}
Solution 2/ Tester's Solution
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int mxN = (int)1e5;
bool prime[mxN+1];
// Create a boolean array "prime[0..n]" and initialize
// all entries it as true. A value in prime[i] will
// finally be false if i is Not a prime, else true.
void SieveOfEratosthenes()
{
memset(prime, true, sizeof(prime));
for (int p=2; p*p<=mxN; p++)
{
// If prime[p] is not changed, then it is a prime
if (prime[p])
{
// Update all multiples of p
for (int i=p*p; i<=mxN; i += p)
prime[i] = false;
}
}
}
void solve(){
int n, k;
cin >> n >> k;
vector<int> arr(n);
for(int i = 0;i < n; ++i) cin >> arr[i];
int sum = 0;// To store the sum for every subarray
int cnt = 0;// To store the final count of subarrays
int avgsum = 0;// To store the average for every subarray
int med = 0;// To store the median value for every subarray
multiset<int> f;// For the initial K elements
for(int i = 0;i < k; ++i){
f.insert(arr[i]);
sum += arr[i];
}
avgsum = sum/k;
int mss = f.size();
if(k%2==0){
mss/=2;
}else{
mss/=2;
mss+=1;
}// To handle odd and even cases for median value
int counter = 0;// To distribute elements into m1 and m2 evenly
auto itr = f.begin();
map<int,int> m1, m2;
// m1 to store the first half of the subarray and m2 to store the second half of the subarray
int s1=0, s2=0;
// To store the sizes(including frequency) of m1 and m2
// s1 will increase or decrease if there is a change in size of m1
// s2 will increase or decrease if there is a change in size of m2
for(;itr != f.end(); itr++){
if(counter==mss) break;
m1[*itr]++;s1++;
counter++;
}
for(;itr != f.end(); itr++){
m2[*itr]++;s2++;
}
// Now first K elements are evenly distributed between m1 and m2
auto last = m1.end();
--last;
med = last->first;// median value will always be the last element of m1
if(avgsum-med>=0 && ((prime[avgsum]==0&&prime[med]==0)||(prime[avgsum]!=0&&prime[med]!=0))){
cnt++;
}// Check the condition
for(int i = 0;i < (n-k); ++i){
int td = arr[i];// Element to be deleted
int ta = arr[i+k];// Element to be added
sum -= td;
sum += ta;
avgsum = sum/k;
// Change in sum and average calculated
// DELETE operation :
if(s1>s2){
// if size(m1) > size(m2) so consider median to be the last value of m1
last = m1.end();
--last;
med = last->first;
if(td<=med){// check if no to be deleted is from m1
m1[td]--;s1--;// decrease frequency in m1
if(m1[td]<=0) m1.erase(td);//erase the number from m1 if its frequency becomes 0
}else{// if number to be deleted is from m2
m2[td]--;s2--;// decrease frequency in m2
if(m2[td]<=0) m2.erase(td);// number erased from m2
// so now to balance the contents in both maps decrease the last element frequency in m1
// by one and add that to m2
m1[med]--;s1--;// decrease frequency in m1
if(m1[med]<=0) m1.erase(med);
m2[med]++;s2++;// increase frequency in m2
}
}else if(s1==s2){
// If size(m1) == size(m2) so consider median to be first value of m2
last = m2.begin();
med = last->first;
if(td>=med){// if no to be deleted is in m2
m2[td]--;s2--;
if(m2[td]<=0) m2.erase(td);// erase the number from m2
}else{// if no is to be deleted from m1
m1[td]--;s1--;
if(m1[td]<=0) m1.erase(td);
// Balance the contents of the maps
// Decrease the first value frequency in m2 by one and add it to m1
m1[med]++;s1++;
m2[med]--;s2--;
if(m2[med]<=0) m2.erase(med);
}
}
// ADD operation :
if(s1==s2){
// If size(m1) == size(m2)
if(s1==0){// If s1 == 0 and s2 == 0
m1[ta]++;s1++;// add element to m1
}else{
// consider median to be first value of m2
last=m2.begin();
med=last->first;
if(ta<=med){// if no is to be added in m1
m1[ta]++;s1++;
}else{//if no is to be added in m2
m2[ta]++;s2++;// add the number
// Balance the contents of the maps
// Decrease the first value frequency in m2 by one and add it to m1
m1[med]++;s1++;
m2[med]--;s2--;
if(m2[med]<=0) m2.erase(med);
}
}
}else if(s1>s2){
// If size(m1) > size(m2) consider median to be last value of m1
last = m1.end();
--last;
med=last->first;
if(ta>med){// if no is to be added in m2
m2[ta]++;s2++;
}else{// if no is to be added in m1
m1[ta]++;s1++;// add the number
// Balance the contents of the maps
// Decrease the frequency of the last value in m1 by one and add it to m2
m2[med]++;s2++;
m1[med]--;s1--;
if(m1[med]<=0) m1.erase(med);
}
}
last = m1.end();
--last;
med = last->first;
// Now map contents are balanced so median == last value of m1
if(avgsum-med>=0 && ((prime[avgsum]==0&&prime[med]==0)||(prime[avgsum]!=0&&prime[med]!=0))){
cnt++;
}// Check the condition
}
cout << cnt << "\n";
}
int32_t main(){
ios::sync_with_stdio(false);
cin.tie(nullptr); cout.tie(nullptr);
SieveOfEratosthenes();
// Use Sieve to find(or precompute) all prime numbers upto 10^5
int tc = 1;
cin >> tc;
while(tc--){
solve();
}
return 0;
}
This problem can also be solved using Merge Sort Tree.
Feel free to share your approach. Suggestions are always welcomed.