친구 (JUNGOL 12491)

https://jungol.co.kr/problem/12491?cursor=MjcsMCwz

---

아이디어

두 포인터를 이용해서 어떤 인덱스를 기준으로 거리가 K1이하인 점들의 범위를 구해줍시다. 그리고 그 안에 있는 모든 학생의 학교별 인원수를 체크해줍니다. 이를 두개의 포인터를 더 사용해서 거리가 K2이하인 점들까지 고려해줍니다. 그러면 각 사람마다 O(1)에 개수를 세어줄 수 있습니다.

---

소스코드

#include <iostream>
#include <vector>
#include <algorithm>
#include <array>
#define ll long long
using namespace std;

ll Count1[505050]={0},Count2[505050]={0};
ll Ans[505050]={0};

int main()
{
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    ll N,K1,K2,i,a,b;
    vector <array <ll,3> > v;

    cin>>N>>K1>>K2;
    for (i=0;i<N;i++)
    {
        cin>>a>>b;
        v.push_back({a,b,i});
    }

    sort (v.begin(),v.end());

    ll L1=0,R1=0,L2=0,R2=0;
    Count1[v[0][1]]++,Count2[v[0][1]]++;
    for (i=0;i<N;i++)
    {
        while (v[i][0]-v[L1][0]>K1) Count1[v[L1++][1]]--;
        while (v[i][0]-v[L2][0]>K2) Count2[v[L2++][1]]--;
        while (R1!=N-1 && v[R1+1][0]-v[i][0]<=K1) Count1[v[++R1][1]]++;
        while (R2!=N-1 && v[R2+1][0]-v[i][0]<=K2) Count2[v[++R2][1]]++;

        Ans[v[i][2]]+=Count1[v[i][1]]-1;
        Ans[v[i][2]]+=(R2-L2+1)-Count2[v[i][1]];
    }

    for (i=0;i<N;i++)
        cout<<Ans[i]<<' ';
}