AVGMED - Editorial

PRACTICE
CONTEST

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.

7 Likes