题解:P17137 [KOI 2026 #1] 朋友

· · 题解

题意简述

n 个点 (x_i,s_i),对于每个点统计满足以下条件之一的点 j(i\not=j) 的数量:

题目分析

考虑容斥,求出 |x_i-x_j|\le k_2 的点的数量,减去 s_i=s_j,|x_i-x_j|\le k_2 的点的数量,加上 s_i=s_j,|x_i-x_j|\le k_1 的点的数量,减去 1(自己)。

那么直接按 x_i 排序后二分就做完了,时间复杂度 O(n\log n)

代码

#include<bits/stdc++.h>
using namespace std;
int n,i,k1,k2;
vector<int>v[500005];
struct T{int x,s,i;}t[500005];
bool operator<(T a,T b){return a.x<b.x;}
int main(){
    cin.tie(0)->sync_with_stdio(0);
    cin>>n>>k1>>k2;
    for(i=1;i<=n;i++){
        cin>>t[i].x>>t[i].s;
        t[i].i=i;
    }
    sort(t+1,t+n+1);
    for(i=1;i<=n;i++)v[t[i].s].push_back(t[i].x);
    for(i=1;i<=n;i++){
        t[i].s=
            (upper_bound(t+1,t+n+1,T{t[i].x+k2,0})
            -lower_bound(t+1,t+n+1,T{t[i].x-k2,0}))
            -(upper_bound(v[t[i].s].begin(),v[t[i].s].end(),t[i].x+k2)
            -lower_bound(v[t[i].s].begin(),v[t[i].s].end(),t[i].x-k2))
            +(upper_bound(v[t[i].s].begin(),v[t[i].s].end(),t[i].x+k1)
            -lower_bound(v[t[i].s].begin(),v[t[i].s].end(),t[i].x-k1))-1;
    }
    sort(t+1,t+n+1,[](T a,T b){return a.i<b.i;});
    for(i=1;i<=n;i++)cout<<t[i].s<<' ';
}