Table of contents

개요

이 자료구조는 point update와 range query를 해야 하는 상황에서 사용된다. 더 자세히 말하자면 길이가 $ n $ 인 수열에 $ a $ 대해

point update: ( $ i $ , $ x $ ) => $ a[i] $ 를 $ x $ 로 치환
range query: ( $ l $ , $ r $ ) => $ a[l] \sim a[r] $ 중 sum/max/min 등 구하기.

와 같은 작업들을 $ q $ 개 해야 한다고 할 때 이를 이진트리를 이용하여 $ O(n + q \log n) $ 안에 처리할 수 있는 자료구조이다. 이 때 sum, max, min 중 어떤 것인지는 트리를 구성하기 전에 정해야 하고 중간에 바꿀 순 없다. 사실 이 세 연산 외에도 교환법칙과 결합법칙이 만족하는, 연산의 순서를 마음대로 바꿔도 되는 모든 연산에 대해서 가능하기 때문에 최대공약수, 최소공배수, 비트연산의 and/or/xor 등 다양한 것들에 적용할 수 있다.

구현

예를 들어 초기 상태의 배열 $ a $ 가 $ [2, 5, 7, 3, 12, 9] $ 이고 range query로 max값을 구할 수 있는 인덱스 트리를 만들어보자.

초기 세팅

우선 이 $ a $ 의 원소를 말단 노드로 가지고, 각 부모들은 두 자식의 값 중 max값을 가지는 완전 이진트리를 구성한다. 이 때 $ a $ 의 길이가 2의 거듭제곱보다 부족한 부분은 원하는 연산에서 항상 버려질 더미 값으로 채우면 된다. 예를 들어 위의 경우 모든 원소가 양수라는 조건 하에 0으로 채우면 된다.

이제 이를 실제 코드에서 구현해보자. 트리의 구현은 인덱스가 1부터 시작하는(1을 루트로 가지는) 배열을 사용해서 구현할 수 있다. 이 배열을 $ bit $ 라고 했을 때 $ bit[x] $ 의 부모는 $ bit[x / 2] $ 가 되고 두 자식은 $ bit[2x] $ , $ bit[2x + 1] $ 로 나타낼 수 있다.

우선 이 트리의 크기를 구해야 하므로 $ n $ 보다 크거나 같은 가장 작은 2의 거듭제곱수를 구해야 한다. 이를 $ base $ 라 하면 다음 코드를 통해 구할 수 있다.

for(base = 1; base < n; base *= 2);

이제 $ bit[base] \sim bit[base + n - 1] $ 에는 원래 배열 $ a $ 의 값이 순서대로 들어가고, $ bit[1] \sim bit[base - 1] $ 에는 각각 자신의 자식들 중 max값을 넣으면 된다.

for(i = 0; i < n; i++) bit[base + i] = a[i];
for(i = base - 1; i > 0; i--) bit[i] = MAXX(bit[2 * i], bit[2 * i + 1]);

여기서 MAXX는 내가 주로 사용하는 max 함수 방식인데, 코드 위에 #define과 삼항 연산자 써서 선언하는 식으로 한다. 이런 방식으로 max나 min을 쓰면 함수를 쓰는 것 보다 훨씬 빨라서 좋다. 다만 이렇게 #define쓰는 것은 컴파일 단계에서 그 부분의 코드를 치환해 버리는 것이기 때문에 주위에 영향을 끼치지 않도록 아래처럼 괄호로 잘 감싸줘야 한다.

#define MAXX(x, y) (((x) > (y)) ? (x) : (y))

여기서 쓰이는 $ bit $ 배열이나 $ base $ , 그리고 아래에서 구현할 함수 등은 class나 struct로 구현해서 내부 변수로 가지고 있어도 되고 전체 코드에서 인덱스트리를 하나밖에 안 쓸 경우 그냥 전역 변수로 선언해서 사용해도 된다.

시간복잡도를 계산해 보면 $ base $ 를 구하는 데 $ O(\log n) $ , 원래 배열의 값을 채우는 데 $ O(n) $ , 나머지 부모 값들을 채우는데 $ O(n) $ 이므로 총 $ O(n) $ 이 걸린다.

point update

이제 point update를 구현해보자. 위에서 봤듯이 이 함수는 인덱스( $ idx $ )와 값( $ x $ )을 받으면 그 인덱스의 값을 갱신하고 이에 따라 트리를 갱신하는 일을 하면 된다.

우선 각 값은 트리에서는 $ base + idx $ 에 저장되므로 이 위치의 값을 갱신해 주면 된다.

idx += base;
bit[idx] = x;

이후 다시 부모가 두 자식의 값 중 max값을 가진다는 성질이 유지되도록 트리를 갱신하면 된다. 하지만 모든 값을 다시 갱신할 필요 없이 어떤 값이 갱신됐을 때 영향을 받는 값은 그 값의 조상들 뿐이라는걸 알 수 있다. 또한 트리의 꼭대기까지 갔다는 판별은 $ idx $ 값이 0인 경우에 루프에서 나오면 되니 아래와 같이 구현하면 된다.

for(idx /= 2; idx; idx /= 2) {
    bit[idx] = MAXX(bit[2 * idx], bit[2 * idx + 1]);
}

따라서 이를 총 정리하면 아래와 같다.

void update(int idx, int x) {
    idx += base;
    bit[idx] = x;
    for(idx /= 2; idx; idx /= 2) {
        bit[idx] = MAXX(bit[2 * idx], bit[2 * idx + 1]);
    }
}

update 함수의 시간복잡도를 구해보면 for문이 반복되는 횟수는 $ O(\log idx) = O(\log n) $ 이다. 또한 for문 안의 내용은 $ O(1) $ 안에 처리가 가능하다. 따라서 총 시간복잡도는 $ O(\log n) $ 이다.

range query

range query에서는 두 인덱스 $ l $ , $ r $ 이 주어지면 그 범위 중 max값을 찾는 작업을 하면 된다. 여기서 우리가 max tree를 구현한 이유가 나온다.

max tree의 말단 노드를 제외한 각 값은 자신의 자식들 중의 max값이기 때문에 어떤 범위를 대표하는 max값이라고 할 수 있다. 예를 들어 위의 트리에서 두 번째 줄의 7은 인덱스가 0~3인 값 중 max값이 7이란 것을 의미하고, 세 번째 줄의 12는 인덱스가 4~5인 값 중 max값이 12란 것을 의미한다. 즉, 우리가 원하는 범위를 정확히 커버하는 대표값들을 찾으면 이 대표값들만 비교해도 그 범위의 max값을 구할 수 있다.

우선 위에서 했듯이 $ l $ , $ r $ 값을 받으면 실제 트리에 적용하기 위해 $ base $ 값을 더해 준다. 그 후 이 범위를 대표할 수 있는 대표값을 찾아보자. 만약에 $ l $ 과 $ r $ 이 같다면 그냥 $ bit[l] $ 이 바로 대표값이 된다. 따라서 우선 $ l $ 과 $ r $ 이 다른 상황에서 살펴보자.

먼저 제일 왼쪽에 있는 값인 $ bit[l] $ 에 대해 생각해보자. 이 값은 $ bit[l / 2] $ 의 왼쪽 자식이거나 오른쪽 자식일 것이다. 만약에 오른쪽 자식일 경우(즉, $ l $ 이 홀수일 경우) $ bit[l / 2] $ 는 이 범위의 대표값에 속할 수 없다. 왜냐하면 $ bit[l / 2] $ 는 $ bit[l - 1] $ 도 같이 대표해 주지만 $ bit[l - 1] $ 은 우리가 구하고자 하는 범위에 속해있지 않기 때문이다. 따라서 이 경우 $ bit[l] $ 이 이 범위를 대표하는 대표값 중 하나라고 빼놓고 $ l $ 을 1 증가시킨다.

제일 오른쪽에 있는 값에 대해서도 같은 작업을 할 수 있다. 만약에 $ bit[r] $ 이 왼쪽 자식일 경우(즉, $ r $ 이 짝수일 경우) $ bit[r / 2] $ 는 이 범위의 대표값에 속할 수 없다. 따라서 이 경우 $ bit[r] $ 이 이 범위를 대표하는 대표값 중 하나라고 빼놓고 $ r $ 을 1 증가시킨다.

이렇게 $ l $ 과 $ r $ 을 상황에 따라 1씩 조정해서 $ l $ 이 왼쪽 자식(짝수), $ r $ 이 오른쪽 자식(홀수)가 되도록 했다고 하자. 그러면 이제 범위 안에 속하는 모든 값들에 대해서 그 부모들은 이 값들을 대표할 수 있다. 따라서 재귀적으로 생각해서 다시 $ l / 2 $ 부터 $ r / 2 $ 의 범위 중 대표값을 구하면 된다.

종료조건은 $ l $ 과 $ r $ 이 같아질 경우, 혹은 $ l $ 과 $ r $ 이 붙어있는데 $ l $ 이 홀수이고 $ l $ 이 짝수라서 1씩 조정하는 과정에서 $ r $ 이 $ l $ 보다 작아지는 경우이다. 하지만 잘 생각해보면 $ l $ 과 $ r $ 이 같은 경우에도 1씩 조정을 한 후 2로 나눠서 다음 단계가 되면 $ r $ 이 $ l $ 보다 작아짐을 알 수 있다. 따라서 $ l \leq r $ 을 만족하는 동안 이 과정을 반복하도록 하면 우리가 원하는 결과를 얻을 수 있다.

int query(int l, int r) {
    l + =base;
    r += base;
    int res = 0;

    while(l <= r) {
        if(l % 2 == 1) {
            res = MAXX(res, bit[l]);
            l++;
        }
        
        if(r % 2 == 0) {
            res = MAXX(res, bit[r]);
            r--;
        }
        
        l /= 2;
        r /= 2;
    }

    return res;
}

query에서의 시간복잡도를 계산해 보면 while문은 $ l $ 과 $ r $ 이 반씩 줄어들면서 적어도 모두 0이 되면 끝나므로 최대 $ O(\log n) $ 번 반복하게 된다. 또한 while문 안에 있는 내용들은 모두 $ O(1) $ 안에 처리가 가능하다. 따라서 총 시간복잡도는 $ O(\log n) $ 이다.

따라서 초기 세팅에 $ O(n) $ , 각 쿼리마다 $ O(\log n) $ 이 걸리므로 총 $ q $ 개의 쿼리를 처리하는 경우 $ O(n + q \log n) $ 이 걸린다. 또한 위의 코드에서는 보기 쉽도록 * 2, / 2 등을 사용했지만 이 연산들은 모두 비트 연산으로 대체할 수 있다. 게다가 짝/홀 판별도 1과 and하는 식으로 판별이 가능하고 $ base $ 를 더하는 것도 $ base $ 가 그 인덱스보다 큰 2의 거듭제곱수임을 생각하면 |연산으로 처리가 가능하다. 즉, 인덱스 트리를 구현하는데 대부분의 연산이 비트 연산으로 가능하기 때문에 이를 활용하면 약간의 성능 향상도 얻을 수 있다.

적용: LIS

인덱스 트리로 풀 수 있는 문제 중 LIS라는 유명한 문제가 있다. LIS는 Longist Increasing Subsequence의 약자로 우리말로 하면 최장 길이 증가 부분수열 이다. 즉, 어떤 수열이 주어졌을 때 여기서 모든 원소가 증가하는 순서를 만족하도록 부분수열을 추출하는데(substring이 아니라 subsequence이므로 연속할 필요는 없다.) 그 중 가장 길이가 긴 것의 길이를 찾는 문제이다.

이 문제는 간단한 DP를 사용하면 $ O(n^2) $ 에 해결할 수 있다. 원래의 수열을 $ a $ 라고 할 때

$ d[i] = a[i]로 ~ 끝나는 ~ 증가 ~ 부분수열 ~ 중 ~ 가장 ~ 긴 ~ 것의 ~ 길이 $

라고 정의하자. 그러면 최종 답은 이 $ d[i] $ 들 중의 최댓값을 구하면 된다.

이제 $ d[i] $ 를 구해보자. $ d[i] $ 는 $ a[i] $ 로 끝나야 하므로 $ a[i] $ 를 어느 수열의 끝에 붙여야 가장 긴 길이를 만들 수 있을지에 대해 생각하면 된다. $ a[i] $ 를 뒤에 붙이기 위해서는 당연히 인덱스가 $ i $ 보다 작아야 하고 또한 그 때의 마지막 수가 $ a[i] $ 보다 작아야 $ a[i] $ 를 붙여도 증가수열임을 유지할 수 있다. 이러한 조건들을 만족하는 부분수열들 중에서 가장 길이가 긴 부분수열의 뒤에 $ a[i] $ 를 붙이면 된다. 즉,

$ d[i] = (j < i, a[j] < a[i] 를 ~ 모두 ~ 만족하는 ~ j ~ 중에 ~ 가장 ~ 큰 ~ d[j]) + 1 $

이라고 할 수 있다. 따라서 반복문으로 $ d[0] \sim d[i-1] $ 을 순회하면서 이 최댓값을 찾으면 각 $ d[i] $ 를 채우는 데 $ O(i) $ 가 걸리므로 총 $ O(n^2) $ 이 걸린다.

이제 이를 인덱스 트리를 통해 최적화를 해보자. 인덱스 트리를 설계할 때는 크게 네 가지를 생각하면 된다.

어떤 종류(min, max, sum 등등)의 트리를 만들 것인가
인덱스는 무엇으로 잡을 것인가
어떤 데이터를 어떤 순서로 언제 어디에 넣을 것인가
쿼리는 언제 어느 범위로 물어볼 것인가

우리의 메인 task는 $ j < i, a[j] < a[i] 를 ~ 모두 ~ 만족하는 ~ j ~ 중에 ~ 가장 ~ 큰 ~ d[j] $ 를 구하는 것이다. 우선 가장 $ d[j] $ 를 구해야 하니 max tree를 구현해야 하고, 트리의 내용은 $ d $ 배열의 값으로 채우면 된다는 것을 짐작할 수 있다. 그러면 자연스럽게 트리의 인덱스는 배열의 인덱스를 그대로 따라가면 된다는 것도 알 수 있다. 이제 $ j < i, a[j] < a[i]를 ~ 모두 ~ 만족 $ 이라는 부분만 해결하면 된다.

먼저 $ j < i $ 부분은 range query를 할 때 트리의 인덱스가 배열의 인덱스이므로 query(0, i-1)을 함으로써 해결할 수 있다. 따라서 $ a[j] < a[i] $ 부분만 해결하면 된다. 이는 트리에 데이터를 넣는 순서로 해결할 수 있다. 트리에 데이터를 넣을 때 $ a[i] $ 가 작은 순으로 넣으면 현재 $ d[i] $ 를 구하려고 할 때 트리에 들어가 있는 값들은 모두 $ a[x] $ 가 $ a[i] $ 보다 작은 값들 이므로 여기서 query(0, i-1)를 구하면 우리가 원하는 $ j < i, a[j] < a[i] 를 ~ 모두 ~ 만족하는 ~ j ~ 중에 ~ 가장 ~ 큰 ~ d[j] $ 를 얻을 수 있다. 이를 정리하면 아래와 같다.

$ a[i] $ 가 작은 순으로 정렬한 뒤 $ a[i] $ 가 작은 $ i $ 부터 시작
초기값은 모두 0으로 채운 max 인덱스 트리 세팅
query(0, i-1)을 구해서 $ j < i, a[j] < a[i] ~ 를 ~ 모두 ~ 만족하는 ~ j ~ 중에 ~ 가장 ~ 큰 ~ d[j] $ 를 얻음
$ d[i] $ 에 위에서 구한 $ d[j] + 1 $ 을 넣고 이 값을 트리에서도 갱신
모든 $ i $ 를 순회하면서 3 ~ 4 반복
$ d[i] $ 들 중 최댓값 출력

이 알고리즘의 시간복잡도를 계산해보자. 우선 $ a $ 를 정렬하는 데 $ O(n \log n) $ 이 걸린다. 이후 인덱스 트리에서는 총 update가 $ n $ 번, query가 $ n $ 번이므로 총 $ O(n \log n) $ 이 걸린다. 그리고 $ d[i] $ 들 중 최댓값을 구하는 것은 한번만 순회하면 되므로 $ O(n) $ 안에 구할 수 있다. 따라서 총 $ O(n \log n) $ 이 걸린다.

이렇게 인덱스 트리는 먼저 DP를 구상한 후 DP에서 구해야 하는 값을 인덱스 트리를 이용해 최적화시키는 용도로도 많이 사용된다.

참고로 LIS는 인덱스 트리를 사용하지 않아도 $ d[i] = 길이가 ~ i인 ~ LIS ~ 중 ~ 가장 ~ 작은 ~ 끝의 ~ 값 $ 으로 정의하면 이분탐색을 통해 $ O(n \log n) $ 안에 해결할 수 있다.

응용1: 트리 펼치기

개요

위에서 보면 인덱스 트리는 보통 배열에 들어가야 할 값들을 말단 노드로 삼아서 구성한다는 것을 알 수 있다. 그렇다면 선형 자료구조가 아닌, 특히 알고리즘에서 자주 나오는 트리구조에 대해서는 인덱스 트리 같은 작업을 하는 것이 불가능할까? 트리가 비선형 구조라서 불가능하다면 강제로 선형 구조로 만들어버리면 된다. 여기서 설명할 트리 펼치기 기법을 사용하면 트리를 선형 자료구조로 만들 수 있다. 이후 여기에서 인덱스 트리를 구성해서 다양한 일들을 할 수 있다.

트리를 펼치는 방법에는 여러가지가 있지만 그 중 하나는 dfs로 순회하면서 방문한 노드를 순서대로 적는 것이다. 예를 들어 위의 트리에 대해 이 방법을 적용하면 $ 1, 2, 4, 2, 5, 2, 6, 8, 6, 9, 6, 2, 1, 3, 7, 3, 1 $ 이라는 수열이 나온다. 이제 선형 자료구조인 수열이 나왔으니 dfs 순회 순서의 의미를 잘 생각하면서 여기에 인덱스 트리를 설계하면 다양한 문제를 해결할 수 있다. 참고로 이렇게 만들어진 수열의 길이는 항상 $ 2n + 1 $ 이다.

적용: LCA

위의 내용만 보면 그래서 뭘 할 수 있을 지 잘 감이 안 올테니 예시를 하나 살펴보자. 이 트리 펼치기 기법과 인덱스 트리를 활용해서 해결할 수 있는 문제 중 LCA라는 문제가 있다. LCA는 Least Common Ancient의 약자로 우리말로 하면 최소 공통 조상이다. 즉, 트리에서 어떤 두 노드가 주어졌을 때 그 두 노드의 공통 조상 중 가장 아래에 있는 조상을 찾는 문제이다.

이 문제를 해결하기 위해 트리 펼치기 기법으로 사용한 dfs 순회 수열과 두 노드의 공통 조상간의 연간 관계를 살펴보자. 5번 노드와 9번 노드 사이의 수열만 따로 추출해보면 $ 5, 2, 6, 8, 6, 9 $ 이다. $ 2 $ 처럼 자식이 있는 경우 원래의 수열에서 여러번 등장하기도 하는데 이런 경우 가장 길게 감싼 부분 수열로 추출하면 된다. 이런 식으로 두 노드를 골라서 그 노드들로 감싸진 부분 수열을 관찰하면 다음 사실을 알 수 있다.

이 부분 수열에는 두 노드의 최소 공통 조상이 포함된다.
이 부분 수열에는 두 노드의 최소 공통 조상을 제외한 공통 조상은 포함되지 않는다.

이 두 관찰 사실을 증명해보자. 두 노드 $ x $ , $ y $ 의 최소 공통 조상을 $ a $ 라고 하자. 우선 트리의 구조 상 $ x $ 에서 $ y $ 로 가는 경로에 $ a $ 는 항상 포함되므로 이 부분수열 안에는 무조건 $ a $ 가 있어야 한다. 또한 dfs로 순회를 하기 때문에 $ a $ 의 조상들이 모두 나온 다음에 $ x $ 와 $ y $ 를 탐색하기 시작할 것이고, 또한 $ a $ 의 자식들을 모두 탐색해야 다시 거슬러 올라가기 때문에 $ x $ 와 $ y $ 를 모두 탐색한 후에야 a의 조상들이 다시 나올 것이다. 따라서 $ a $ 의 조상은 이 부분 수열에 포함될 수 없다.

이를 이용해 문제를 풀기 위해 트리의 노드들에 번호를 다시 붙이는 작업을 하는데, dfs 순으로 해도 되고 bfs 순으로 해도 되고 아니면 자신이 좋아하는 순서로 해도 되지만 부모는 자식들보다 항상 숫자가 작다는 조건은 만족해야 한다. 예를 들어 위에서 예시로 든 트리는 bfs 순서대로 번호를 매겼기 때문에 이 조건을 아주 잘 만족한다.

이제 이 조건과 위의 두 관찰에 의해 두 노드로 인해 만들어진 부분 수열에서 가장 작은 수가 최소 공통 조상의 번호임을 알 수 있다. 먼저 트리 펼치기를 할 때 각 노드별로 자신이 처음 나오는 인덱스와 마지막에 나오는 인덱스를 저장한다. 그리고 펼쳐진 수열로 min 인덱스 트리를 구성한다. 이후 x, y의 공통 조상을 구하라는 쿼리가 오면

query(min(x의 첫 인덱스, y의 첫 인덱스), max(x의 마지막 인덱스, y의 마지막 인덱스))

를 구하면 $ O(\log n) $ 만에 구할 수 있다. 따라서 트리의 노드 숫자가 $ n $ , 공통 조상을 물어보는 쿼리의 수가 $ q $ 일 경우 $ O(n + q \log n) $ 에 해결할 수 있다.

응용2: 압축

위에서는 인덱스로 잡은 변수들이 모두 보기 좋게 $ 0 \sim n $ 까지의 값이었다. 하지만 상황에 따라서 $ 0 \sim 10^9 $ 과 같이 큰 범위에서 주어진 n개의 값들을 인덱스로 삼고 싶은 하는 경우도 생긴다. 그렇다고 크기가 $ 10^9 $ 짜리 트리를 만들면 시간이던 공간이던 터져버릴 것이다. 이럴 때 쓰는 기술이 압축이다.

압축을 쓰기 위한 전제 조건은 $ n $ 개의 수들이 대소관계 외에 각 값들에는 의미가 없어야 한다. 이러한 경우에 각 수를 정렬한 뒤 크기 관계를 유지하도록 $ 0 \sim n-1 $ 까지 수를 새로 부여한 다음 이를 토대로 인덱스 트리를 구성하면 된다.

적용: 다시 한번 LIS

위에서 한 번 살펴봤던 LIS를 다시 한번 살펴보자. 우리가 얻은 DP식은 $ d[i] = (j < i, a[j] < a[i] 를 ~ 모두 ~ 만족하는 ~ j ~ 중에 ~ 가장 ~ 큰 ~ d[j]) + 1 $ 였다. 여기서 $ j < i, a[j] < a[i] $ 라는 두 가지 조건을 만족하는 값을 찾아야 했고, $ j < i $ 는 쿼리의 인덱스 조건으로, $ a[j] < a[i] $ 는 트리에 넣는 순서로 해결했다. 하지만 잘 생각해보면 둘 다 똑같이 하나가 다른 하나보다 작다는 크기 비교의 연산인데 꼭 $ i $ 값을 인덱스로 하고 $ a[i] $ 값을 순서로 해야 할 필요가 있을까? 당연히 없다. 즉, $ a[j] < a[i] $ 를 쿼리의 인덱스 조건으로 하고 $ j < i $ 를 트리에 넣는 순서로 해결해도 LIS를 구할 수 있다.

여기서 $ a[j] < a[i] $ 를 인덱스로 쓰기 위해서 해야 하는 작업이 압축이다. LIS의 경우에 increasing 조건을 위해서는 수의 대소관계만 중요하지 수 자체의 값이 결과에 영향을 주지 않는다. 즉, 만약에 원래 배열의 값이 $ [2, 5, 7, 3, 12, 9] $ 였다면 이를 $ [0, 2, 3, 1, 5, 4] $ 로 바꿔도 결과에는 영향을 주지 않는다.

따라서 아래와 같이 하면 $ a[j] < a[i] $ 를 쿼리의 인덱스 조건으로 하고 $ j < i $ 를 트리에 넣는 순서로 해결하는 LIS를 구현할 수 있다.

$ a[i] $ 가 작은 순으로 정렬한 뒤 $ 0 \sim n-1 $ 까지의 값을 부여 (압축), 이 값을 $ p[i] $ 라 하자
초기값은 모두 0으로 채운 max 인덱스 트리 세팅
$ i = 0 $ 부터 시작
query(0, p[i]-1)을 구해서 $ p[j] < p[i], j < i 를 ~ 모두 ~ 만족하는 ~ j ~ 중에 ~ 가장 ~ 큰 ~ d[p[j]] $ 를 얻음
$ d[p[i]] $ 에 위에서 구한 $ d[p[j]] + 1 $ 을 넣고 이 값을 트리에서도 갱신
모든 $ i $ 를 순회하면서 4 ~ 5 반복
$ d[i] $ 들 중 최댓값 출력

응용3: 두 가지 조건은 무조건 인덱스 트리?

위의 LIS에서 보면 $ j < i, a[j] < a[i] $ 라는 두 개의 대소관계 조건을 해결해야 하는 문제를 하나를 쿼리의 인덱스 조건, 다른 하나를 트리에 넣는 순서로 해서 해결했다. 그렇다면 이를 확장해서 이와 같은 생각을 해볼 수 있지 않을까?

두 가지 대소관계 조건을 만족해야 하는 문제는 LIS를 이용해서 하나는 쿼리의 인덱스 조건, 다른 하나는 트리에 넣는 순서로 구현하면 해결할 수 있다!

그리고 이는 실제로 가능하다. 위와 같은 상황을 최대한 일반화 해보자.

두 집합 $ P, Q $ 가 있을 때
$ Q $ 의 임의의 원소 $ q $ 에 대해
$ f_1 (p) \lt f_2 (q) $ 와
$ f_3 (q) \lt f_4 (p) \lt f_5(q) $ 를 모두 만족하는 모든 $ p $ 에 대해
$ f_6(p) $ 중 sum/max/min 등을 구해야 하는 경우 인덱스 트리로 해결이 가능하다.

예를 들어 위의 LIS의 경우는

$ P = Q = \{ 0, 1, \ldots , n-1 \} $
$ f_1 (x) = f_2 (x) = x $
$ f_3 (x) = -1 , f_4 (x) = f_5 (x) = a[x] $
$ f_6 (x) = d[x] $

라고 볼 수 있다.

이는 아래와 같이 구현하면 된다.

$ f_3 (q), f_4 (p), f_5(q) $ 의 후보들을 압축
1에서 $ f_3 (q) $ 이 부여 받은 값을 $ s[q] $ , $ f_5(q) $ 이 부여 받은 값을 $ e[q] $ , $ f_4 (p) $ 이 부여 받은 값을 $ i[p] $ 라 하자.
$ f_3 (q), f_4 (p), f_5(q) $ 의 후보의 양 만큼의 크기로 인덱스 트리 세팅한다.
$ f_2 (q) $ 가 작은 $ q $ 부터 $ f_2 (q) $ 순으로 시작
지금의 $ f_2 (q) $ 보다 $ f_1 (p) $ 이 작고 아직 트리에 들어가지 않은 모든 $ p $ 에 대해 $ f_6(p) $ 의 값을 $ i[p] $ 의 인덱스에 넣고 갱신
5의 과정은 $ f_6(p) $ 의 값에 대해 미리 정렬을 해두면 모든 과정을 통틀어서 $ P $ 의 크기만큼 밖에 걸리지 않음
query(s[q] + 1, e[q] - 1)을 구함
모든 $ q $ 를 순회하면서 5 ~ 7 반복

이 알고리즘의 정당성을 살펴보자. 우선 우리는 $ q $ 를 $ f_2 (q) $ 가 작은 순서로 시작한다. 또한 쿼리를 구하기 전에 $ f_2 (q) $ 보다 $ f_1 (p) $ 이 작고 아직 트리에 들어가지 않은 모든 $ p $ 에 대해 트리를 갱신한다. 따라서 지금 트리에는 $ f_1 (p) \lt f_2 (q) $ 를 만족하는 모든 값이 들어가있다.

이제 query(s[q] + 1, e[q] - 1)를 구하면, 이는 압축 과정을 따랐으므로 여기서 구해진 값은 모두 $ f_3 (q) \lt f_4 (p) \lt f_5(q) $ 를 만족한다.

따라서 $ f_1 (p) \lt f_2 (q) $ 와 $ f_3 (q) \lt f_4 (p) \lt f_5(q) $ 라는 두 가지 조건을 모두 만족하는 것에 대한 값을 구해야 하는 문제를 해결할 수 있다.

일반화를 하기 위해 수식으로 번잡하게 썼지만 실제로 구현을 해보면 생각보다 복잡하지 않다.

적용: Concatenation with intersection

최근에 codeforces에서 풀었던 문제 중 Concatenation with intersection라는 문제가 이렇게 두 가지 대소관계 조건을 만족해야 하는 상황을 인덱스 트리를 사용해서 잘 해결했던 경험이 있어서 소개하고자 한다. 참고로 이 문제는 당시 참가자들 중 1.5%만 해결했던 문제다.

우선 문제 설명부터 간단히 하자면, 세 문자열 $ a, b, s $ 가 주어진다. $ a, b $ 의 길이 $ n $ 은 50만 이하이며 두 문자열의 길이는 같고, $ s $ 의 길이 $ m $ 은 $ 2n $ 이하이다.

$ a, b $ 에서 각각 연속한 부분 문자열을 추출한 다음 $ a $ 에서 추출한 문자열 뒤에 $ b $ 에서 추출한 문자열을 붙여서 새로운 문자열을 생성한다. 이 때 다음 두 가지 조건을 만족하도록 부분 문자열을 추출하는 경우의 수를 세면 된다.

이렇게 생성한 문자열이 $ s $ 와 동일해야 한다.
$ a $ 에서 추출한 위치와 $ b $ 에서 추출한 위치의 교점이 있어야 한다. 즉, $ a[x] $ 는 $ a $ 에서 추출한 문자열에 속하고 $ b[x] $ 는 $ b $ 에서 추출한 문자열에 속하는 $ x $ 가 적어도 하나는 있어야 한다.

인덱스 트리를 사용하려면 위에 처럼 애매한 조건이 아니라 크기 관계가 되는 조건으로 바꿔야 한다. 이를 위해 몇 가지 전처리를 해야 한다.

우선 $ a $ 의 각 위치 $ i $ 에 대해 $ s $ 의 시작을 이 위치에 맞췄을 때 $ s $ 의 앞에서부터 최대 몇 글자까지 일치하는지에 대한 값들을 $ A[i] $ 에 저장해둔다. 이는 z-알고리즘을 사용하면 $ O(n+m) $ 안에 구할 수 있다.

$ b $ 에 대해서는 각 위치 $ i $ 에 대해 $ s $ 의 끝을 이 위치에 맞췄을 때 $ s $ 의 뒤에서부터 최대 몇 글자까지 일치하는지에 대한 값들을 $ B[i] $ 에 저장해둔다. 이는 $ s $ 와 $ b $ 를 뒤집은 문자열에 대해 z-알고리즘을 사용하고 다시 뒤집으면 구할 수 있다.

이제 어떤 위치 $ x, y $ 에 대해 $ A[x] $ 와 $ B[y] $ 의 의미를 그림으로 나타내면 아래와 같다.

이 상황에서 $ A[x] + B[y] $ 가 $ m $ 보다 크거나 같다면 우리는 $ x $ 에서 부터 앞에서 $ l $ 만큼 추출하고 $ y $ 에서 부터 뒤에서 $ m - l $ 만큼 추출하면 이를 붙인 문자열은 $ s $ 와 같게 된다. 따라서 이 경우에 $ A[x] + B[y] - m + 1 $ 의 경우의 수 만큼 $ s $ 를 만들 수 있다. 또한 여기서 위에서 언급된 교점의 조건까지 생각해보자. 교점이 생기려면 $ x $ 가 $ y $ 보다 작거나 같아야 하고 둘의 길이를 합쳐서 $ m $ 만큼 잘랐을 때 가운데에 교점이 생길 만큼 $ x $ 와 $ y $ 가 가까워야 한다. 이를 식으로 나타내면 아래와 같다.

$ A[x] + B[y] >= m $
$ x <= y $
$ y - x + 1 <= m - 1 $

이를 다시 정리하면 아래와 같이 쓸 수 있다.

$ m - A[x] - 1 < B[y] $
$ y - m + 1 < x < y + 1 $

와! 우리가 원하는 두 가지의 크기관계 조건이 나왔다.

이제 마지막으로 우리가 구해야 하는 값에 대해 생각해보자. 우리는 모든 $ y $ 에 대해서, 위를 만족하는 모든 $ x $ 에 대해 $ A[x] + B[y] - m + 1 $ 의 값을 구한 뒤 다 더하면 된다. 이는 위를 만족하는 $ x $ 의 갯수와 $ A[x] $ 의 합을 통해 구할 수 있다. 더 자세히 말하자면, 이를 만족하는 $ x $ 의 갯수를 $ t $ , 이에 대한 $ A[x] $ 의 합을 $ SA $ 라 하면 $ SA + t * (B[y] - m + 1) $ 이 우리가 원하는 값이 됨을 알 수 있다. 또한 이 두 값은 트리의 말단 노드에 $ (A[x], 1) $ 라는 pair를 갱신하는 sum 인덱스 트리를 구현하면 한번의 쿼리로 구할 수 있다.

이를 토대로 인덱스 트리는 아래와 같이 구현할 수 있다.

$ res = 0 $ 으로 초기화한다.
인덱스는 각 $ x $ 나 $ y $ 의 값에 $ m $ 을 더한 값을 사용한다. ( $ y - m + 1 $ 이 양수가 되게 하기 위함)
트리의 크기는 $ n + m + 1 $ 이 된다.
$ -A[x] $ 를 기준으로도 미리 정렬을 해둔다. (트리에 갱신하는 과정을 총 $ O(n) $ 안에 하기 위함)
$ B[y] $ 를 기준으로 정렬해서 작은 순으로 $ y $ 를 순회한다.
현재의 $ y $ 에 대해 $ m - A[x] - 1 < B[y] $ 를 만족하면서 아직 트리에 들어가지 않은 $ x $ 의 $ (A[x], 1) $ 를 $ m + x $ 의 인덱스에 갱신한다.
query(m + y - m + 2, m + y) = query(y + 2, m + y)를 구한다.
이 결과를 $ (SA, t) $ 라 하면 $ SA + t * (B[y] - m + 1) $ 를 $ res $ 에 더해준다.
모든 $ y $ 를 순회하면서 5 ~ 7 반복
$ res $ 출력

z-배열 알고리즘의 시간복잡도는 $ O(n+m) $ 이고 인덱스 트리의 시간복잡도는 $ O(n + m + n \log (n + m)) $ 이다. 따라서 $ m <= 2n $ 임을 생각하면 총 시간복잡도 $ O(n \log n) $ 안에 위 문제를 해결할 수 있다.

인덱스 트리 (Indexed Tree)

개요

구현