Solution #
Iterate over each frequency. Suppose now we are on frequency . Put all stations with frequency in the vector and all radio stations with frequency into the vector.
Now we want to calculate the number of pairs such that the left radio station is from the vector and the right station is from vector.
Sort the vector by position and sort the vector by the left bound of the stations’ range. Iterator the stations in the vector and put all the stations in the vector which can reach the current station in the axis(actually we need to put them in some data structure). Now we need to know how many stations in the axis can be reached by the current station. This can be done with some range-sum-query data structure(like fenwick tree): we add one on the position for each new station and use range query to find the stations we want. However, since the positions are up to we also need to compress the coordinate, which is really annoying, so a simpler way to do this is to use a balanced BST in pb_ds library to find the order directly.
The lesson learnt is that when we want to find the order, especially with coordinate compression, consider pb_ds.
Code #
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include<ext/pb_ds/priority_queue.hpp>
using namespace __gnu_pbds;
#define forn(i, n) for (int i = 0; i < int(n); ++i)
#define for1(i, n) for (int i = 1; i <= int(n); ++i)
#define ms(a, x) memset(a, x, sizeof(a))
#define F first
#define S second
#define all(x) (x).begin(),(x).end()
#define pb push_back
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
const int INF = 0x3f3f3f3f;
mt19937 gen(chrono::steady_clock::now().time_since_epoch().count());
template<typename... T> void rd(T&... args) {((cin>>args), ...);}
template<typename... T> void wr(T... args) {((cout<<args<<" "), ...);cout<<endl;}
using ordered_set = tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n,k;
cin>>n>>k;
vector<pii> fre[10005];
forn(i,n){
int x,r,f;
cin>>x>>r>>f;
fre[f].pb({x,r});
}
ll ans=0;
auto solve=[&](vector<pii>& left,vector<pii>& right){
sort(all(left));
sort(all(right),[](pii a,pii b){return a.F-a.S<b.F-b.S;});
ll res=0;
int i=0;
ordered_set tree;
for(auto it:left){
while(i<right.size()&&right[i].F-right[i].S<=it.F){
tree.insert(right[i].F);
i++;
}
res+=tree.order_of_key(it.F+it.S+1)-tree.order_of_key(it.F+1);
}
return res;
};
for(int i=1;i<=1e4;i++){
if(fre[i].empty()) continue;
vector<pii> left(all(fre[i])),right;
for(int j=max(1,i-k);j<=i+k&&j<=10000;j++){
right.insert(right.end(),all(fre[j]));
}
ans+=solve(left,right);
}
cout<<ans;
return 0;
}