Tutorial for Codeforces 762E - Radio stations

Solution #

Iterate over each frequency. Suppose now we are on frequency ii. Put all stations with frequency ii in the leftleft vector and all radio stations with frequency [ik,i+k][i-k,i+k] into the rightright vector.

Now we want to calculate the number of pairs such that the left radio station is from the leftleft vector and the right station is from rightright vector.

Sort the leftleft vector by position and sort the rightright vector by the left bound of the stations’ range. Iterator the stations in the leftleft vector and put all the stations in the rightright 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 10910^9 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;
}