#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;
}
Update : On changing the custom comparator function from cmp to less, the solution passed.
Custom comparator function :
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?
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.
Yeah, I understand that something is wrong with my comparator function. But can you explain why the above didn't work?
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.