BOJ 7469 K번째 수 (in Q log N!)

https://www.acmicpc.net/problem/7469

7469번: K번째 수

Step 1 ($O(M \log^3 N)$)

이 단계는 머지 소트 트리를 알고 있다면 누구나 풀 수 있습니다. 머지 소트 트리는 일반적인 세그먼트 트리와 비슷한데, $[l, r]$ 구간을 담당하는 노드에는 $[l, r]$ 구간에 있는 수들을 정렬해서 배열 형태로 저장해 둔 세그먼트 트리를 말합니다. 머지 소트 트리의 구축은 모든 노드에 대해 정렬을 독립적으로 하면 $O(N \log^2 N)$에 가능합니다. 정렬을 독립적으로 하지 않고, 머지 소트에서 두 배열을 합치는 방식으로 구현하면 시간 복잡도를 $O(N \log N)$으로 줄일 수 있습니다.

이제 어떻게 쿼리를 처리하는지 알아봅시다. $[l, r]$ 구간에서 $K$번째로 작은 수를 구해 봅시다. 편의상 L은 1부터 $N$까지의 정수로 이루어진 순열이라고 가정합시다. (순열이 아니라면 기초적인 좌표 압축을 이용해 순열로 만들어 줄 수 있습니다.)

바로 문제를 해결하기 쉽지 않기 때문에, 답을 $x$라고 두고 이분 탐색을 해서, $[l, r]$ 구간에서 $x$보다 작은 수가 $K$개 이하인지 초과인지를 알아보겠습니다. 일반적인 세그먼트 트리와 같이, $[l, r]$ 구간은 최대 $O(\log N)$개의 노드의 합집합으로 나타낼 수 있습니다. 머지 소트 트리의 각 노드에서 이분 탐색을 해서 $[l, r]$ 구간에서 $x$ 이하의 수가 몇 개인지를 찾아줄 수 있습니다. 이제 각 노드에 대해 이분 탐색으로 얻은 결과를 합친 것이 $[l, r]$ 구간에서 $K$ 이하인 수의 개수가 됩니다.

시간 복잡도를 살펴봅시다. 일단 답 $x$에 대해 이분 탐색을 하기 때문에 로그가 하나 붙고, 각 노드별로 이분 탐색을 하므로 로그가 하나 더 붙고, 그렇게 이분 탐색을 하는 노드도 로그 개이기 때문에 로그는 총 세 개가 붙습니다. 따라서 이때의 시간 복잡도는 $O(M \log^3 N)$이고, 이것만으로도 문제를 해결하기에는 충분합니다. (사실 이게 가장 잘 알려진 풀이입니다.)

Step 2 ($O(M \log^2 N)$)

여기서 로그를 하나 뗄 수 있습니다. 로그를 하나 떼기 위해서는 아이디어가 필요합니다. 아무래도 이분 탐색을 두 번 중첩해서 하는 것은 낭비 같습니다. 일반적으로 세그먼트 트리 문제에서 파라메트릭 서치를 할 때는, 흔히 "세그 이분 탐색"이라고 불리는 아이디어로 로그를 하나 떼는 것이 가능합니다. 이 문제에서도 이 트릭을 활용할 수 있을까요?

직접 이 트릭을 적용하기는 힘듭니다. 하지만 여기서 아이디어를 하나 낼 수 있습니다. 이제부터는 수열 A에서 생각하지 말고, 수열 A의 역순열 B에서 문제를 해결해 봅시다. ($B[A[i]] = i$와 같이 정의됩니다.) 이때, 문제는 B에서 $l$ 이상 $r$ 이하의 원소만 남겨놓았을 때, 이들의 위치 중 $K$번째로 작은 것을 찾는 것으로 바뀝니다.

이제 여기까지 왔으면, 머지 소트 트리 위에서 이분 탐색을 함으로써 이를 찾아줄 수 있습니다. 자세한 과정은 다음과 같습니다. 노드 $x$에서 $[s, e]$ 구간 안에 들어가는 $K$번째 원소를 구한다고 해 봅시다. 먼저 $x$의 왼쪽 자식에서 이분 탐색을 해서, $[s, e]$ 안에 들어가는 원소가 몇 개인지를 셉니다. 만약 이것이 $K$개 이상이라면, 답은 왼쪽 노드에 있으니 왼쪽 노드로 들어가 계속 탐색합니다. 만약 이것이 $K$개 미만이라면, 답은 오른쪽 노드에 있으니 $K$에서 앞서 구한 값을 빼고 오른쪽 노드로 들어가 계속 탐색합니다. 이때 $O(\log N)$개의 노드에서 이분 탐색을 하기 때문에, 시간 복잡도는 $O(M \log^2 N)$입니다.

Step 3 ($O(M \log N)$)

여기서 더 시간 복잡도를 줄일 수는 없을까요? 놀랍게도 가능합니다. 핵심은 머지 소트 트리의 각 노드에서 이분 탐색을 하지 않는 것입니다. 어떻게 머지 소트 트리에서 이분 탐색을 하지 않냐고요? 이를 위해서는 머지 소트 트리에 전처리를 해 두어야 합니다.

설명을 간편하게 하기 위해, 위에서 설명했떤 로그제곱 최적화는 잠시 잊고 처음으로 돌아가 보겠습니다. 우리는 로그 세제곱 풀이에서 $O(l, r, s, e)$와 같은 함수를 구현해야 하는데, 이 함수는 배열의 $[l, r]$ 구간에서 $s$ 이상 $e$ 이하인 수의 개수를 찾는 함수입니다. 원래 배열이 순열이라고 가정했으므로, $1 \le s \le e \le N$입니다.

머지 소트 트리의 첫 번째 노드에서 이분 탐색을 하지 않고 $s$ 이상 $e$ 이하인 수의 개수를 찾는 방법은 무엇일까요? 머지 소트 트리의 루트에서는 $1$ 이상 $N$ 이하의 모든 원소가 존재할 테니, $e-s+1$이 답이 됩니다. 그렇다면 루트 노드의 자식 노드는 어떨까요? 이 경우에도 비슷한 방법을 적용할 수 있을까요?

이를 간편하게 처리하기 위해, 머지 소트 트리의 모든 노드에 좌표 압축 (Coordinate Compression 또는 Renumbering) 을 시행합니다. 예를 들어, $A=[1, 5, 8, 2, 4, 3, 7, 6]$, $(l, r) = (1, 4)$, $(s, e) = (3, 6)$라고 가정합시다. 머지 소트 트리의 루트의 왼쪽 자식 노드를 생각합시다. 이 노드가 관리하는 구간에 있는 수의 종류는 $[1, 5, 6, 2]$입니다. 따라서 이들을 $[1, 4]$ 구간에 좌표 압축하면, $[1, 5, 6, 2] \rightarrow [1, 3, 4, 2]$로 변하게 됩니다.

좌표 압축을 했다면, 머지 소트 트리의 노드를 타고 내려가는 과정에서 $(s, e)$를 좌표 압축된 형태로 변형해 주어야 합니다. 예를 들어, 위의 예시에서 처음에 $(s, e) = (3, 6)$이었다면 좌표 압축된 형태로 변형한 이후 $(s, e)$는 $(3, 4)$가 됩니다. 이렇게 $(s, e)$를 바꾸고 나면 처음에 루트 노드에서 했던 것처럼, 단순히 $e-s+1$을 계산하는 것으로 노드 하나에서 쿼리를 마칠 수 있습니다.

이제 남은 과제는 하나입니다. $(s, e)$를 어떻게 빠르게 바꿀 것인지를 알아내야 합니다. 이것을 빠르게 하지 못하면, 원래처럼 이분 탐색을 이용해 $(s, e)$를 계속 변형해야 하고, 그러면 말짱 도루묵이니까요. 이제 이것을 빠르게 계산하기 위해 적당한 전처리가 필요합니다. 각 노드별로 $s$ 값을 왼쪽 노드로 옮길 때 어떤 값으로 전이시켜야 하는지를 전처리할 수 있다고 합시다. 이를 오른쪽 노드에 대해서도 똑같은 방식으로 계산하고, $e$에 대해서도 똑같은 방식으로 처리하면 결국 문제가 해결됩니다.

이 방법은 단순합니다. 머지 소트 트리의 노드 $x$에 대해, 노드 $x$와 $x$의 왼쪽 자식 $x \rightarrow l$에 들어 있는 두 배열에서 투 포인터를 시행합니다. 이때 투 포인터를 통해서 $x$에서 $s$였던 값이 $x \rightarrow l$에서 얼마로 전이되어야 하는지의 값을 $x$에 존재하는 모든 값에 대해 각각 구해줄 수 있습니다. 이를 $x \rightarrow r$에 대해서도 똑같이 하면 문제가 해결됩니다.

설명이 다소 난해하므로, 코드를 첨부합니다. 아래 코드는 144ms로 돌고, 세제곱 코드에 비해 큰 차이가 없지만, 벡터 접근이 많기 때문에 벡터 대신 동적 할당을 이용해 코딩하면 훨씬 나은 성능을 보일 수 있을 것입니다.

 

#include <bits/stdc++.h>

typedef long long ll;

using namespace std;

vector<int> tree[400002];
vector<int> lPnt[200002], rPnt[200002];

void init(int i, int l, int r, int *A){
    if(l==r){
        tree[i].push_back(A[l]);
        return;
    }
    int m = (l+r)>>1;
    init(i*2, l, m, A);
    init(i*2+1, m+1, r, A);
    tree[i].resize(r-l+1);
    lPnt[i].resize(r-l+1);
    rPnt[i].resize(r-l+1);
    merge(tree[i*2].begin(), tree[i*2].end(), tree[i*2+1].begin(), tree[i*2+1].end(), tree[i].begin());
    for(int x=0, y=-1, z=-1; x<r-l+1; x++){
        if(y < m-l && tree[i*2][y+1] == tree[i][x]) y++;
        if(z < r-m-1 && tree[i*2+1][z+1] == tree[i][x]) z++;
        lPnt[i][x] = y, rPnt[i][x] = z;
    }
}

int query(int i, int l, int r, int s, int e, int x){ /// x번째로 작은 값을 찾는다.
//        printf("Query %d %d %d %d %d %d\n", i, l, r, s, e, x);
    if(l==r) return l;
    int m = (l+r)>>1;
    int tmp = lPnt[i][e] - (s==0 ? -1 : lPnt[i][s-1]);
    if(x <= tmp) return query(i*2, l, m, (s==0 ? -1 : lPnt[i][s-1]) + 1, lPnt[i][e], x);
    else return query(i*2+1, m+1, r, (s==0 ? -1 : rPnt[i][s-1]) + 1, rPnt[i][e], x-tmp);
}

int n, q;
int arr[100002], idx[100002], loc[100002];
vector<int> renumber;

int main(){
    scanf("%d %d", &n, &q);
    for(int i=0; i<n; i++) scanf("%d", &arr[i]);

    vector<pair<int, int> > vec;
    for(int i=0; i<n; i++) vec.push_back(make_pair(arr[i], i));
    sort(vec.begin(), vec.end());
    for(int i=0; i<n; i++) idx[vec[i].second] = i;
    for(int i=0; i<n; i++) loc[idx[i]] = i;
    init(1, 0, n-1, loc);

    while(q--){
        int l, r, x;
        scanf("%d %d %d", &l, &r, &x);
        int ret = query(1,  0, n-1, l-1, r-1, x);
        printf("%d\n", vec[ret].first);
    }
}