수열과 쿼리 8 [BOJ 13553]

Mo's Algorithm과 세그먼트 트리가 완전히 체화되어 있다면 쉽게 풀 수 있는 문제이지만, 그게 아니라면 아이디어를 떠올리기 조금 힘들다. 그래서 별 다른 아이디어 없이도 P1의 티어를 받은 것 같다.

 

문제 요약

길이 10만 이하의 수열 A에서 다음 쿼리를 수행한다.

l r : 구간 [l, r]에 있는 i, j(i<=j)에 대해서 abs(A[i]-A[j])<=K인 쌍의 수를 출력한다.

 

사용 알고리즘

Mo's Algorithm

펜윅 트리

 

풀이

만약 어떤 구간 [s, e]의 원소들이 세그먼트 트리에 표시되어 있다면, 우리는 어떤 수 x와 차이가 K 이하인 원소의 수를 그 세그먼트 트리에 쿼리를 하나 날려줌으로써 구할 수 있다. - (*)

그렇다면 한 쿼리에서 다른 쿼리로 이동하는 것을 시간제한 안에 할 수 있어야하는데, 해당 쿼리들은 업데이트 쿼리가 없는 오프라인 쿼리들이므로 Mo's Algorithm을 활용하여 정렬하고, 상태를 이동할 때마다 O(logN)의 시간복잡도로 세그먼트 트리를 갱신해줄 수 있다. 이 과정에서 (*)을 통해 문제에서 구하고자 하는 값의 변화를 실시간으로 계산해줄 수 있다. 총 시간복잡도는 O(NsqrtNlogN)이 된다.

시간복잡도에 비해 시간제한이 조금 빡세기 때문에 필자는 펜윅 트리로 구현하였다. 그랬는데도 1904ms가 나오는 것을 보면 아마 세그먼트 트리를 사용하면 시간초과가 날 것 같다.

 

코드

#include <bits/stdc++.h>
#define SIZE 100002
typedef long long ll;
using namespace std;
struct Query
{
    int idx, s, e;
};
int A[SIZE], tree[8*SIZE], sqrtN, L=200000, K;
ll ans[SIZE], res;
Query q[SIZE];
bool comp(Query x, Query y)
{
    if (x.s/sqrtN==y.s/sqrtN) {
        if (x.e==y.e) return x.s<y.s;
        return x.e<y.e;
    }
    return x.s/sqrtN<y.s/sqrtN;
}
void add(int pos, int v)
{
    while (pos<=L) tree[pos]+=v, pos+=(pos&-pos);
}
int get(int pos)
{
    int ret=0;
    while (pos>0) ret+=tree[pos], pos-=(pos&-pos);
    return ret;
}
void mo_add(int x)
{
    res+=get(A[x]+K)-get(A[x]-K-1);
    add(A[x], 1);
}
void mo_dis(int x)
{
    add(A[x], -1);
    res-=get(A[x]+K)-get(A[x]-K-1);
}
int main()
{
    int N, M, i;
    scanf ("%d %d", &N, &K);
    sqrtN=sqrt(N);
    for (i=1; i<=N; i++) scanf ("%d", &A[i]);
    scanf ("%d", &M);
    for (i=0; i<M; i++) {
        scanf ("%d %d", &q[i].s, &q[i].e);
        q[i].idx=i;
    }
    sort(q, q+M, comp);
    int s=q[0].s, e=q[0].e;
    for (i=s; i<=e; i++) mo_add(i);
    ans[q[0].idx]=res;
    for (i=1; i<M; i++) {
        while (s<q[i].s) mo_dis(s++);
        while (s>q[i].s) mo_add(--s);
        while (e<q[i].e) mo_add(++e);
        while (e>q[i].e) mo_dis(e--);
        ans[q[i].idx]=res;
    }
    for (i=0; i<M; i++) printf("%lld\n", ans[i]);
}