수열과 쿼리 23 [BOJ 16979]

수열과 쿼리 8, 9와 함께 풀면 좋은 문제이다. 아직 안 풀었다면 함께 풀어보는 것을 추천한다.

 

문제 요약

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

s e : 구간 [s, e]에 존재하는 i<j에 대해 Ai>Aj인 (i, j) 쌍의 수를 출력

 

사용 알고리즘

펜윅 트리

Mo's Algorithm

 

풀이

어떤 수열에 대해 쿼리 1 N을 처리해보자.

이를 처리하는 것은 어렵지 않다. 처음 원소부터 하나씩 세그먼트 트리에 넣으면서, 이미 들어간 원소 중 자신보다 큰 원소의 수를 쿼리로 세주면 된다. 시간복잡도는 O(NlogN)이다.

이를 통해 우리는 쿼리를 처리하기 위해서 어떤 구간의 모든 수를 담고 있는 세그먼트 트리가 있어야한다는 사실을 알 수 있다. 그렇다면 이를 Mo's Algorithm으로 쿼리와 쿼리 사이 전이를 하며 업데이트를 함으로써 해결할 수 있고, 왼쪽에 업데이트 되는 것과 오른쪽에 업데이트 되는 것을 적절히 처리해주면 된다.

어떤 수열의 왼쪽에 수가 추가된다면, 이미 세그먼트 트리에 있는 수들 중 자신보다 작은 수의 개수를 더해주면 된다.

반대로 오른쪽에 수가 추가된다면, 이미 세그먼트 트리에 있는 수들 중 자신보다 큰 수의 개수를 더해주면 된다.

따라서 O((NsqrtN+M)logN)의 시간복잡도에 문제를 해결할 수 있다. N의 범위가 10만으로 작음에도 불구하고 시간제한이 5초나 된다는 것은 단순 O(NlogN)이 아니라 거기에 sqrtN이나 logN이 더 붙는다는 것을 짐작할 수 있고, 이도 문제를 풀 때 길잡이가 되는 요소 중 하나이니 주의 깊게 봐야한다. 또한 이렇게 시간복잡도가 큰 경우에는 상수를 줄여야할 필요가 있을 때도 있어서, 필자의 코드는 세그먼트 트리 대신 펜윅 트리를 사용하였다.

 

코드

#include <bits/stdc++.h>
#define SIZE 100002
typedef long long ll;
using namespace std;
struct Query
{
    int idx, s, e;
};
int A[SIZE], tree[SIZE], k, L;
pair<int, int> v[SIZE];
ll ans[SIZE], res;
Query q[SIZE];
bool comp(Query x, Query y)
{
    if (x.s/k==y.s/k) {
        if (x.e==y.e) return x.s<y.s;
        return x.e<y.e;
    }
    return x.s/k<y.s/k;
}
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_radd(int x)
{
    res+=get(L)-get(A[x]);
    add(A[x], 1);
}
void mo_rdis(int x)
{
    res-=get(L)-get(A[x]);
    add(A[x], -1);
}
void mo_ladd(int x)
{
    res+=get(A[x]-1);
    add(A[x], 1);
}
void mo_ldis(int x)
{
    res-=get(A[x]-1);
    add(A[x], -1);
}
int main()
{
    int N, M, i;
    scanf ("%d %d", &N, &M);
    k=sqrt(N);
    for (i=1; i<=N; i++) {
        scanf ("%d", &v[i].first);
        v[i].second=i;
    }
    sort(v+1, v+N+1);
    int x=1;
    A[v[1].second]=x;
    for (i=2; i<=N; i++) {
        if (v[i].first!=v[i-1].first) x++;
        A[v[i].second]=x;
    }
    L=x;
    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_radd(i);
    ans[q[0].idx]=res;
    for (i=1; i<M; i++) {
        while (s<q[i].s) mo_ldis(s++);
        while (s>q[i].s) mo_ladd(--s);
        while (e<q[i].e) mo_radd(++e);
        while (e>q[i].e) mo_rdis(e--);
        ans[q[i].idx]=res;
    }
    for (i=0; i<M; i++) printf("%lld\n", ans[i]);
}