Am I using policy based ordered_set wrong? Please help

#	User	Rating
1	ecnerwala	3648
2	Benq	3580
3	orzdevinwang	3570
4	cnnfls_csy	3569
5	Geothermal	3568
6	tourist	3565
7	maroonrk	3530
8	Radewoosh	3520
9	Um_nik	3481
10	jiangly	3467

#	User	Contrib.
1	maomao90	174
2	adamant	164
2	awoo	164
4	TheScrasse	160
5	nor	159
6	maroonrk	156
7	-is-this-fft-	150
8	SecondThread	147
9	orz	146
10	pajenegod	145

Question link : https://www.codechef.com/problems/AVGMED Question summary : Find number of subarrays of length K, such that average >= median and also the median and average must both be prime or non prime.

Solution link using ordered_set : https://www.codechef.com/viewsolution/39036972 Am I doing something wrong with how to use ordered_set? Or my solution is wrong itself? Please help.

Code

#include<bits/stdc++.h>
using namespace std;

#define LL long long
#define FI ios_base::sync_with_stdio(false); cin.tie(NULL);
#define PREC cout << setprecision(10) << fixed;

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

template <class T> struct cmp {
	bool operator() (T x, T y) {
		return x.first < y.first;
	}
};
template <class T> using ordered_set = tree<T, null_type, cmp<T>, rb_tree_tag, tree_order_statistics_node_update>;

#define int long long

const int MOD = 1e9 + 7;
const int INF = 1e9 + 9;
const int MX = 1e5 + 5;

bool isPrime[MX];
void sieve() {
	memset(isPrime, true, sizeof isPrime);
	isPrime[0] = isPrime[1] = false;
	for(int i=2;i*i<MX;i++) {
		if(isPrime[i] == false) continue;
		for(int j=i*i;j<MX;j+=i) {
			isPrime[j] = false;		
		}
	} 
}

int n, k;
int inp[MX];

#undef int 
int main() {
	#define int long long
	FI;
	sieve();
	int tt;
	cin >> tt;
	while(tt--) {
		cin >> n >> k;
		for(int i=0;i<n;i++) {
			cin >> inp[i];
		}

		ordered_set<pair<int,int> > st;
		int ans = 0;

		int sum = 0;
		for(int i=0;i<k;i++) {
			st.insert({inp[i],i});
			sum += inp[i];
		}
		int avg = sum/k;
		auto it = st.find_by_order((k-1) /2);
		assert(it != st.end());
		int med = it->first;

		assert(avg >= 0 && med >= 0);
		if(isPrime[avg] == isPrime[med] && avg >= med) ans++;

		for(int i=k;i<n;i++) {
			sum += inp[i] - inp[i-k];
			st.erase({inp[i-k], i-k});
			st.insert({inp[i], i});

			avg = sum/k;
			it = st.find_by_order((k-1) /2);
			assert(it != st.end());
			med = it->first;

			assert(avg >= 0 && med >= 0);
			if(isPrime[avg] == isPrime[med] && avg >= med) ans++;
		}
		cout << ans << endl;
	}
	
	return 0;
}

Another solution : https://www.codechef.com/viewsolution/38989875 Similar solution which got accepted. I am not to able to see what I did wrong (I am blind or what?). Can someone kindly help?

Comments (6)

Write comment?

yesnomaybe

3 years ago, # |

Update : On changing the custom comparator function from cmp to less, the solution passed.

Custom comparator function :

template <class T> struct cmp {
	bool operator() (T x, T y) {
		return x.first < y.first;
	}
};

I am trying to create an ordered multiset, and for that using pair. I want this set to be ordered on first value. But I don't see what's wrong with my comparator function though which led to wrong answer?

Is it because there could be multiple entries with same first value but different second value, and due to that the ordered_set is not able to handle "find_by_order" correctly. But I don't see how that should be the case? As I care about order imposed only by first key of the pair.

Can someone please explain what I am missing here? And how to use ordered_set over pair<int,int> i.e ordered_multiset?

→ Reply

ishuu7

I think you've written wrong comparator for the ordered_set. I changed it to this and It worked.

You can check AC code from here.

3 years ago, # ^ |

Yeah, I understand that something is wrong with my comparator function. But can you explain why the above didn't work?

the_hyp0cr1t3

Try

template <class T> struct cmp {
	bool operator() (T x, T y) {
		return x.first == y.first? x.second < y.second : x.first < y.first;
	}
};

← Rev. 2 →

Yeah I see that it will work, but can you explain why we need to explicitly write condition for second value in pair? As in I only want set to be sorted as per first value, and in my comparator function if first values are equal then it will return false. As I don't care about the order of entries that have same first value.

Your comparator must ensure strict weak ordering. See link. This is not limited to pbds ordered set. You can replicate similar 'bugs' with STL set.

yesnomaybe's blog