# AVGMED - Editorial

**URL:** <https://discuss.codechef.com/t/avgmed-editorial/79772>\
**Category:** editorial\
**Tags:** c-plus-plus, editorial, merge-sort-tree, policy-based-data-structure, sieve, logi2020\
**Created:** [October 17, 2020, 11:29am UTC](https://discuss.codechef.com/t/avgmed-editorial/79772 "2020-10-17T11:29:11Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![tejasdeore7](https://www.codechef.com/sites/all/themes/abessive/images/user_default_thumb.jpg) [@tejasdeore7](https://discuss.codechef.com/u/tejasdeore7)\
**Post date:** [October 17, 2020, 11:29am UTC](https://discuss.codechef.com/t/avgmed-editorial/79772/1 "2020-10-17T11:29:11Z")

</div>

[PRACTICE](https://www.codechef.com/problems/AVGMED)  
[CONTEST](https://www.codechef.com/LOGI2020/problems/AVGMED)

Author : [tejasdeore7](https://www.codechef.com/users/tejasdeore7)  
Tester : [adarsh\_sinhg](https://www.codechef.com/users/adarsh_sinhg)  
Editorialist : [tejasdeore7](https://www.codechef.com/users/tejasdeore7)

DIFFICULTY :  
Medium

PREREQUISITES :  
[Policy based data structures](https://www.geeksforgeeks.org/policy-based-data-structures-g/)

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](https://www.geeksforgeeks.org/sieve-of-eratosthenes/). 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](https://www.geeksforgeeks.org/ordered-set-gnu-c-pbds/) 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](https://www.geeksforgeeks.org/merge-sort-tree-for-range-order-statistics/).

Feel free to share your approach. Suggestions are always welcomed.

---

<div class="post-metadata">

**Author:** ![mahajanvv](https://discuss.codechef.com/user_avatar/discuss.codechef.com/mahajanvv/32/28539_2.png) [@mahajanvv](https://discuss.codechef.com/u/mahajanvv)\
**Post date:** [March 3, 2021, 10:36am UTC](https://discuss.codechef.com/t/avgmed-editorial/79772/2 "2021-03-03T10:36:11Z")

</div>

When I am trying to submit the following code which is similar to editorial’s submission. I am getting WA.  
can somebody help me out where I am wrong?  
When I am trying to submit editorial’s solution it is getting Accepted but following code is not.

#include\<bits/stdc++.h\>  
#include\<ext/pb\_ds/assoc\_container.hpp\>  
#include\<ext/pb\_ds/tree\_policy.hpp\>  
using namespace \_\_gnu\_pbds;  
using namespace std;  
#define ordered\_set tree\<int, null\_type,less, rb\_tree\_tag,tree\_order\_statistics\_node\_update\>  
#define int long long  
#define test() int t;cin\>\>t;while(t–)  
const int mxN = (int)1e5;  
bool isPrime[mxN+1];  
void calcPrimeNumbers(){  
memset(isPrime, true, sizeof(isPrime));  
for(int i=2; i_i\<=mxN; i++){  
if(isPrime[i] == true){  
for(int j=i_i; j\<=mxN; j+=i){  
isPrime[j] = false;  
}  
}  
}  
}  
int32\_t main(){  
ios::sync\_with\_stdio(false);  
cin.tie(NULL);  
calcPrimeNumbers();  
test()  
{  
ordered\_set s;  
int n,k;  
cin\>\>n\>\>k;  
vector v(n);  
int sumi = 0;  
int cnt = 0;  
for(int i=0; i\< (int)n; i++){  
cin\>\>v[i];  
}  
for(int i=0; i\< (int)k; i++){  
s.insert(v[i]);  
sumi += v[i];  
}  
int avgSum = sumi/k;  
int med = \*s.find\_by\_order((k+1)/2 - 1);  
// bool condition = (isPrime[median] ^ isPrime[avg]);  
if(avgSum-med\>=0 && ((isPrime[med]==0 && isPrime[avgSum]==0)||(isPrime[med]!=0 && isPrime[avgSum]!=0))){  
cnt++;  
}  
for(int i=0; i\<(int)(n-k); i++){  
s.erase(s.find\_by\_order(s.order\_of\_key(v[i])));  
s.insert(v[i+k]);  
sumi -= v[i];  
sumi += v[i+k];  
avgSum = sumi/k;  
med = \*s.find\_by\_order((k+1)/2 - 1);  
// bool condition = (isPrime[median] ^ isPrime[avg]);  
if(avgSum-med\>=0 && ((isPrime[med]==0 && isPrime[avgSum]==0)||(isPrime[med]!=0 && isPrime[avgSum]!=0))){  
cnt++;  
}  
}  
cout\<\<cnt\<\<"\n";  
}  
return 0;  
}