세그먼트 트리 2 (Segment Tree)

세그먼트 트리에 대해서 이전에 부분합을 구하는 것으로만 사용하는 것으로 작성한 적이 있다.

그러나 최근 BOJ에서 1168번 문제와 12899번 문제를 풀면서 다른 활용법을 알게 되었는데

이에 대한 내용이다.

---

 세그먼트 트리의 특성은 간단하게 부모 노드가 모든 자식노드의 합이 된다는 것이다.

위에서 작성한 1168번 문제와 12899번 문제는 이 특성을 이용해 순서를 찾는데 이용하는 문제이다.

 자, 다음과 같은 height 3에 모든 자식 노드가 1의 값을 가지는 세그먼트 트리가 있다.

root node의 값을 통해 우리는 전체 자식노드 중 값을 가지는 노드가 총 4개임을 알 수 있다.

그리고 각 자식노드는 왼쪽부터 1, 2, 3, 4로, 오름차순으로 d=1의 등차수열에 매칭된다.

tree[4] -> 1, tree[5] -> 2... 로 매칭된다는 의미다.

 

이 세그먼트 트리에 있는 수 중에서 3번째로 작은 수를 찾으려고 할 때를 생각해보자.

---

1. tree[1] ( root )의 값은 4이다. 즉, 트리에 수는 총 4개가 있다. 찾는 수는 3번째로 작은 수다.

 - 왼쪽 자식 (tree[2])의 값은 2이다. 왼쪽으로 내려가면 수가 최대 2개이므로 3번째로 작은수는 찾을 수 없다.

  <왼쪽 자식은 1번째 작은 수, 2번째 작은 수 를 가진다.>

 - 오른쪽 자식 (tree[3])의 값은 2이다. 자식노드의 매칭값은 오름차순으로 정렬되어있다.

  즉, 왼쪽 자식 노드의 값만큼의 작은 수를 스킵하고 온 것이므로, 오른쪽 자식 노드는 3번째 작은 수부터 2개를 가진다.

 => 3번째로 작은 수는 오른쪽 자식 밑에서 첫 번째로 작은 수다. ( '원래 찾으려던 순서' - '왼쪽 노드의 수' )

 

 2. tree[3] 의 값은 2이다. 찾는 수는 1번째로 작은 수다.

 - 왼쪽 자식 (tree[6])의 값은 1이다. 왼쪽으로 내려가면 수가 최대 1개이므로, 1번째로 작은 수를 찾을 수 있다.

 => 3번째로 작은 수는 tree[6]에 매칭된 수이다.

 

3. tree[6]에 매칭된 수는 6 - (non-leap 노드의 수) 이므로, 3번째로 작은 수는 3이다.

---

위의 단계와 같이 트리를 이용한 Binary Search 같은 방식으로 진행된다.

 

BOJ 12899. 데이터 구조의 문제를 이용해 한번 더 생각해보자.

찾아야 될 값의 범위는 [1, 2_000_000]이다.

즉, 구성할 세그먼트 트리의 단말노드의 수가 최소 200만개가 있어야 한다는 의미이다.

200만보다 높은 가장 작은 2의 지수승은 2^21 = 2,097,152이므로, 트리의 크기는 2^22로 설정한다.

 

먼저 유형 1. 트리에 자연수 X를 추가한다.

이것은, Bottom-Up 방식으로, 먼저 X 자리에 추가한 뒤 부모 노드를 찾아가면서 1씩 증가시켜주면된다.

 - X에 매칭되는 노드는 '2^21 + (X - 1)' 이다. (단말노드의 시작이 tree[2^21]이고, X의 범위는 1부터 이므로 1을 빼준다.)

 - 노드의 key를 2로 나눠서 부모노드를 찾아가면서 1씩 더해준다.

이를 Java코드로 표현하면 다음과 같다.

public void set(int[] tree, int X){
    int key = (1 << 21) + (X - 1);
    while(key > 0) {
        tree[key]++;
        key /= 2;
    }
}

 

 그 다음으로 유형 2. 트리에 포함된 수 중 X번째로 작은 수를 출력하고 삭제한다.

이 문제에서는 트리에 X개 이상의 원소가 있음이 보장되지만, 없을 경우를 요구하는 문제가 있다면,

tree[1]과 비교하여 'X > tree[1]'일 경우 출력이 불가능함을 표시할 수 있다.

 - X번째로 작은수를 찾는 방법은 위쪽에서 설명한 바와 같이 Top-Down 방식으로

   root노드부터 찾는 X번쨰 수가 왼쪽에 있는지 오른쪽에 있는지를 비교하면서 내려가면된다.

 - 찾으면서 내려가는 노드들의 값을 1씩 감소시켜줘야한다. ( 자식노드를 삭제하므로, 개수를 1 감소 시켜야함 )

 

이를 Java코드로 표현하면 다음과 같다.

public int poll(int[] tree, int X) {
    int node = 1, nonLeap = (1 << 21);
    while(node < nonLeap) {
        tree[node]--; // 개수 1 감소
        int left = node * 2, right = left + 1;
        if(tree[left] >= X) { // 왼쪽과 오른쪽 자식 비교
            node = left;
        } else {
            node = right;
            X -= tree[left];
        }
    }
    tree[node]--;
    return (node - nonLeap) + 1;
}

 

이제 이 두 함수를 이용해 유형 1,2를 수행하고 출력하는 코드만 작성하면 된다.

 

 

---

 

 다음으로, 이런 성질을 응용한 BOJ 1168. 요세푸스 문제 2 다.

바뀐 점은

 1. 트리에 값을 입력할 필요 없이, N개의 자식노드가 모두 값을 가진 상태로 시작한다.

 2. X번째 작은 수가 아니라, 이전에 탐색한 수 부터 K번째 다음 수를 찾아야 한다.

으로 두가지로 생각할 수 있다.

 

먼저 트리를 구성해야하는데, 트리의 크기는 N개의 자식노드가 필요하므로,

2^(Math.ceil( log_2(N) ) + 1) 만큼의 크기를 가지면된다.

예제 입력을 예로 들면, N=7이므로 log_2(7) = 2.8, 올림하면 3, +1하면 4, 2^4 = 16의 크기가 필요하다.

자바에는 log_2가 없으므로, log_10(N) / log_10(2)를 통해 값을 얻는다.

이를 Java코드로 나타내면 다음과 같다.

public int estimateTreeSize(int N) {
    return (1 << ((int) Math.ceil(Math.log(N) / Math.log(2)) + 1));
}

 

자식 노드에 모두 1을 설정해야하므로, tree.size / 2 번째 노드부터 모두 1씩 설정하고, 부모까지 적용한다.

public int[] init(int size, int N) {
    int[] tree = new int[size];
    for(int i = 0 ; i < N ; i++) { // 단말 노드의 1번째부터 N번째 까지 적용한다.
        int key = (size / 2) + i;
        while(key > 0) { // 부모 노드에도 개수를 반영한다.
            tree[key]++;
            key /= 2;
        }
    }
    return tree;
}

 

 이제 바뀐 점에서의 1번은 모두 해결했으니 2번을 해결해야한다.

이를 수행하는 과정은 다음과 같다.

 1. 현재 위치의 왼쪽에 몇개의 수가 있는지 확인한다.

 2. 왼쪽에 있는 수 + K 번째 수를 찾는다.

 3. 오른쪽에 있는 수가 부족할 경우, 오른쪽에 남은 수만큼 빼고 1번째부터 다시 탐색한다.

 

먼저 1번은 현재 노드의 위치부터 위로 거슬러 올라가면서,

자신이 올라온 노드가 오른쪽 자식일 경우, 왼쪽 자식의 값만큼 카운터를 추가하면된다.

예를 들어, 가장 먼저 3번째 수를 제거한 후, 위치를 찾는 법을 그림으로 보면 다음과 같다.

이를 java 코드로 보면 아래와 같다.

public int getOrder(int[] tree, int loc) {
    int cnt = 0;
    while(loc > 0) {
        if(loc % 2 == 1) cnt += tree[(loc / 2) * 2];
        loc /= 2;
    }
    return cnt;
}

 

그리고 2번 과정은 getOrder로 구한 값+K를 이용해 위에서 다른 문제를 풀었던 방식으로 찾으면 된다.

 

이제 3번 과정인데, 먼저 오른쪽에 수가 부족한지는 getOrder()+K가 tree[1]보다 큰지를 통해 알 수 있다.

예시를 이용해 예를 들면, 3, 6이 지워진 다음, getOrder()+K를 수행한 값을 알아보자.

 getOrder()의 값은 4가 반환되고, 4+K = 7인데, tree[1]은 5이므로 7-5를 해서 2번째 값을 찾으면 된다.

만약, 남은 노드의 수가 적고, K가 클 경우 위의 계산을 한 번만 하면 범위 안으로 들어오지 못할 수 있으므로,

while를 이용해 tree[1] >= order+K가 될 때 까지 계속 계산한다.

public int checkOrderOverflow(int[] tree, int order, int K) {
    int find = order+K;
    while(tree[1] < find) find -= tree[1];
    return find;
}

 

그리고 이 함수들을 적절히 이용하면 문제를 해결할 수 있다.