자료구조 - 세그먼트 트리(Segment Tree)

 

1. 세그먼트 트리(Segment Tree, 구간 트리)란?

 

특정 구간 내 연산(쿼리)에 대해 빠르게 응답하기 위해 만들어진 자료구조이다.

예를 들어 크기가 N=100인 int배열 arr이 있다면 1~100의 인덱스 내 숫자들이 위치해 있을 것이다.( 0을 사용하지 않는다고 가정 )

이 때, 이 배열의 구간 arr[l] ~ arr[r]의 합을 구하고자 한다고 하자. 아래의 그림과 같다.
(1 <= l, r <= 100 and l < r)

 

구간 합 배열 예시

 

물론 간단히 반복문을 사용해서 l의 위치와 r의 위치를 찾아서 덧셈을 수행하는 것도 방법이기는 하다.
그와 같이 연산 시, 전체 시간 복잡도는 최악의 경우, O(N)으로 표현할 수 있을 것이다. 이 행위를 M번 수행한다면 O(N*M)이 된다.

그런데, 여기에 연산을 추가하여 배열 내 특정 위치의 값을 다른 값으로 변경한다고 생각해보자. 

예를 들어 arr[p] (단, l <= p, p <= r)이 기존 데이터를 A -> B로 변경한다면 이 작업을 수행하는 행위는 O(1)이 된다. 값이 변경되는 행위를 총 M번 수행한다면 변경에 따른 시간 복잡도는 O(M)이 된다.

즉, 현재 연산은 합 / 변경의 2가지이며 이 2가지 연산을 1회로 보고 최대 M번 수행한다고 고려한다면 총 시간복잡도는 아래와 같다.
총 시간복잡도 : O(N * M) + O(M) = O(N * M)

 

이는 너무나 큰 시간을 소요할 수 있으니 다음과 같이 합 배열(sum)을 따로 생성해 미리 전처리 해두면 어떨까?

인덱스 (i)의 값	1	2	3	4	5
합 (sum)	1	3	6	10	15

 

위와 같이 숫자가 들어오면 그 숫자에 대한 합을 미리 다른 배열에 저장해두는 것이다.

1) 만약 2~3번째 수까지의 합만 구하고 싶다면?
sum[3] - sum[1] = 6 -1 = 5

2) 만약 3 ~ 5 번째 수까지의 합만 구하고 싶다면?
sum[5] - sum[2] = 15 - 3 = 12

이와 같이 진행한다면 합엽산 시, 상수 시간대로 활용하여 진행할 수 있게 된다. 여기에 최초 sum 배열을 전처리하는 시간이 추가된다.
따라서 N개의 데이터에 대해 M회 수행 시 O(N + M) 이 된다. ( 전처리 O(N) / 합 연산 O(M))

그런데, 데이터가 변경된다면? 만약 위의 배열에서 3번째 숫자가 3 → 6으로 변했다고 가정하자. 그러면 sum 배열에서도 3번째 값만 변경하면 되는 것이 아니라 4, 5번째의 합 결과 또한 변경해야 한다. 따라서 수가 바뀌면 M회 진행 시, O(N * M)의 연산이 추가로 발생한다.

따라서 총 시간복잡도는 아래와 같다.
총 시간복잡도 : O(N+M) + O(N*M) = O(N*M)

 

이 때, N과 M이 충분히 큰 숫자라고 한다면 짧은 시간 안에 문제를 풀 수 없게 되는 문제가 발생한다. 세그먼트 트리는 이 연산의 시간 복잡도를 O(MlogN) 만에 수행할 수 있도록 만들어준다.

 

 

2. 세그먼트 트리(Segment Tree)의 구조

 

기본적으로 이진 트리의 구조를 갖는다. 즉, 리프(Leaf) 노드가 아닌 모든 노드는 1개 이상, 2개 이하의 자식(Child) 노드를 보유한다. 그런데 기존 이진 트리에서는 각각의 모든 노드가 고유의 값을 가졌다면 이번에는 부모(Parent) 노드가 자식(Child) 노드들의 합을 저장하는 방식이라고 생각하면 된다.

루트(Root) 노드의 번호는 1번으로 시작하며, 현재 노드의 번호가 X라면 왼쪽 자식(Left Child) 노드의 번호는 2*X, 오른쪽 자식(Right Child) 노드의 번호는 2*X+1 이다.

배열로 구현하기 때문에 1번으로 시작하여 수행하는 것이 일반적이다. 그래야 자식 노드의 번호를 더 직관적으로 구분할 수 있다.(트리 구조에 대한 포스팅을 보면 더 쉽게 이해할 수 있다.)

이 트리는 구간의 합을 구하기 위함이므로 전체 구간의 크기는 리프(Leaf) 노드의 개수와 같다.

예를 들어, 크기가 n인 배열의 구간합을 위한 세그먼트 트리를 구현 시 n=7 이라고 가정한다면 세그먼트 트리는 다음과 같다.

 

세그먼트 트리, 7개 데이터 구간 연산 예시

 

여기서 눈치 챌 수 있는 것은 이진 트리 구조이므로 자식 단계로 깊이가 깊어질수록 해당 층위의 노드는 위 층위의 2배이다. 따라서, 만약 크기 n이 2의 제곱꼴이라면 완전/포화 이진 트리가 될 수 있다.

위의 경우에는 크기가 7이라서 마지막 14, 15번 노드는 빈 값으로 유지하게 되는 것이다.
즉, 데이터의 개수에 따라 저장되는 층위가 다를 수 있다.(하나만 저장하면 되기 때문에 아래 층위로 가지 않음)
전체 필요 노드의 수는 2의 제곱꼴인 경우 2*n -1 개이다.(n이 2의 제곱꼴이 아니라면 불 필요한 노드는 비워둔다)

 

트리의 높이는 어떻게 구할까? 높이는 루트(Root) 노드에서의 가장 긴 경로를 의미한다. 위의 그림에서 루트(Root) 노드부터 리프(Leaf) 노드의 가장 긴 경로는 3이다.

따라서, 만약 n=8이었다면 n이 2의 3제곱이 8이므로 3이 높이이다. 즉, n이 2의 제곱꼴인 경우, 높이 h=logn 이 된다.
위 그림은 n=7이므로 높이는 동일하게 3이지만 log7은 2.xxxx 일 것이다. 즉, 2의 제곱꼴이 아니라면 높이 h=[logn]+1이다.

 

이 세그먼트 트리(Segment Tree)를 배열로 구현한다고 할 때, 필요한 배열의 크기는 어떻게 될까?
완전, 포화 이진 트리라는 가정 하에 필요 배열의 크기는 2^(h+1) - 1이 된다. (n이 2의 제곱꼴이 아니면 남은 위치는 비운다.)

 

3. 세그먼트 트리(Segment Tree) 구현하기

 

1) 생성자

생성자에서는 세그먼트 트리를 배열로 구현하기 위해 필요한 배열의 크기를 설정한다. 아래의 코드를 확인하자.

class SegmentTree{
    long tree[];
    int treeSize;

	public SegmentTree(int arrSize){
        // Tree의 높이는 전체 배열 개수를 log화 한 것
        // Leaf는 n개이고, 전체 개수는 각 leaf를 더한 개수도 포함해야 하므로, log(n)의 높이로 구해야함(소수점이 있어 더 필요하니 ceil 사용)
        // Java의 log는 자연로그 함수이므로 log(2)를 통해 나누어 log2로 나누는 효과를 구함
        int h = (int) Math.ceil(Math.log(arrSize) / Math.log(2));

        // Tree에 들어가는 Node의 개수는 2의 높이+1 제곱의 미만 개이다.
        this.treeSize = (int) Math.pow(2, h+1);
        tree = new long[treeSize];
    }
}

 

자바의 Math.log 함수는 자연로그 이기 때문에 밑을 2로 나누는 효과를 위해 log(2)로 한 번 더 나누어 준다. 위에서 설명하였듯이 이 트리의 높이는 전달되는 배열의 size에 영향을 받는다.

Math.ceil 함수는 올림을 하는 함수이므로, 2의 제곱꼴이 아닌 경우에는 +1을 해주기 위해 설정한 것이다.

 

2) init 메소드 (최초 세그먼트 트리를 구성하기)

이 메소드는 세그먼트 트리의 리프(leaf) 노드에 배열로 전달된 값들을 넣고 그 부모(parent) 노드들에 각각의 구간 합을 삽입하는 방식으로 진행된다. 재귀 방식으로 간단하게 구현할 수 있다.

public long init(int[] nums, int node_idx, int start, int end){
    // 만약 start == end라면 leaf 노드라는 의미
    // 즉, 배열의 값을 그대로 저장함
    if(start == end){
        return tree[node_idx] = nums[start];
    }

    // 위에서 return 못했다면 start != end
    // 그렇다면 좌측 노드와 우측 노드의 합으로 구해짐.
    // 좌측 노드의 idx는 *2 이며 우측은 *2+1 이 된다.
    return tree[node_idx] = init(nums, node_idx*2, start, (start+end)/2) +
        init(nums, node_idx*2+1, (start+end)/2+1 ,end);
}

 

위에서 node_idx는 1로 전달해준다. 왜냐면 트리(Tree) 구조를 배열로 구현하였기 때문에 1에서부터 시작하여 다음 노드의 인덱스 값을 *2 (좌측 자식) / *2+1 (우측 자식)으로 찾아가기 위함이다.

start, end 는 start는 1에서 시작할 것이고 end는 주어진 상황하에서 마지막에 해당하는 값을 넣어주면 된다. 만약 숫자가 총 100개라면 start = 1, end = 100으로 주어지면 된다.

start == end라면 리프(leaf) 노드에 도달했다는 의미이므로 배열의 값을 전달하고 반환하며 그 부모(parent)들은 리프(leaf)에서부터 값이 계속 전달되면서 합으로 정해지는 것을 확인할 수 있다.

 

3) update 메소드

update 메소드는 중간에 값이 변경되었을 경우에 재귀 방식으로 기존의 값과 현재 값의 차이만큼을 영향 받는 모든 노드에 연산하여 변경해주는 방식이다.

public void update(int node_idx, int start, int end, int idx, long diff){
    // 만약 변경할 index 값이 범위 밖이면 해당 tree는 확인 불 필요
    if(idx < start || end < idx) return;

    // 변경된 차이만큼을 영향 받는 각 node에 더해줌
    tree[node_idx] += diff;

    // leaf 노드에 다다르기 까지 모든 노드의 값을 바꿔야 하므로 지속 진행
    if(start != end){
        update(node_idx*2, start, (start+end)/2, idx, diff); // 좌측 node로
        update(node_idx*2+1, (start+end)/2+1, end, idx, diff); // 우측 node로
    }
}

 

여기서도 node_idx는 1로 전달해주고 start / end는 숫자 배열의 시작부터 끝을 전달한다. 즉, start =1 / end = n(n개의 숫자가 주어질 시)으로 전달하고 idx는 변경할 노드의 인덱스 값을 전달하면 된다.

맨 위의 조건문을 통해 영향을 받는 노드인지 아닌지를 확인할 수 있으며 중요한 것은 위 메소드에서는 기존 값과의 차이를 전달했다는 것이다. 그 차이는 미리 연산해서 전달하였다(반드시 그런 것은 아니고 위에선 그렇게 구현함).

이 차이를 전달 시에, 트리를 구현 시에 사용한 배열이 아닌 원래 배열도 값을 당연히 바꾸어 줘야 한다. 왜냐면 동일한 위치의 값을 또 다른 값으로 바꿀 때, 차이가 제대로 전달되지 않을 것이기 때문이다.

 

4) sum 메소드

이는 구간합을 구하기 위한 메소드이다. 일단 코드를 보자.

// 구하고자 하는 덧셈을 진행
public long sum(int node_idx, int start, int end, int left, int right){
    // 범위를 벗어나게 되는 경우 더할 필요 없음
    if(left > end || right < start){
        return 0;
    }

    // 범위 내 완전히 포함 시에는 더 내려가지 않고 바로 리턴
    if(left <= start && end <= right){
        return tree[node_idx];
    }

    // 그 외의 경우 좌 / 우측으로 지속 탐색 수행
    return sum(node_idx*2, start, (start+end)/2, left, right)+
            sum(node_idx*2+1, (start+end)/2+1, end, left, right);
}

 

여기서도 패러미터(Parameter)는 동일하게 전달한다. start / end는 주어진 기존 배열의 처음과 끝이므로 1 / n을 전달한다. left / right는 전체 범위 내의 구간 합을 구할 구간의 시작 / 끝 지점의 값을 전달하면 된다.

구간의 합을 구할 때는 4가지의 경우가 있을 수 있다.

① [left,right]와 [start,end]가 겹치지 않는 경우
  - 겹치지 않으면 해당 구간은 불필요

② [left,right]가 [start,end]와 완전히 일치하는 경우
  - 현재 그 구간합을 보유한 노드가 정답

③ [start,end]가 [left,right]를 완전히 포함하거나 한 쪽만 겹치는 경우
  - 이 때는 left / right의 정확한 구간보다 더 넓은 범위이므로 더 아래로 내려가서 찾아야 한다.

 

다시 이 그림을 보자.

arr[5]~arr[6] 즉, 6번 노드의 결과를 요구하고자 한다고 가정하자.

 

①의 경우 : 루트(root) - Left 탐색 - 2번 노드를 현재 탐색중인 경우
  - 이 경우 2번 노드는 물론 그 아래도 전혀 필요가 없다. 탐색을 중지한다.

②의 경우 : 루트(root) - Right 탐색 - Left 탐색 - 6번 노드 현재 탐색중인 경우
  - 6번 노드가 완전히 포함하므로 더 이상 탐색하지 않고 6번 노드의 결과를 반환

③의 경우 : 루트(root) 즉, 1번 노드 혹은 루트(root) - Right 탐색 - 3번 노드 탐색중인 경우
  - 현재 탐색 범위가 찾고자 하는 범위보다 더 넓으니 더 아래로 탐색을 진행

 

위 각각의 경우가 어떤 경우일지는 위의 트리 그림을 통해 쉽게 이해할 수 있을 것이다.

 

전체 코드는 아래를 확인하자.

package com.test;

class SegmentTree{
    long tree[];
    int treeSize;

    public SegmentTree(int arrSize){
        // Tree의 높이는 전체 배열 개수를 log화 한 것
        // Leaf는 n개이고, 전체 개수는 각 leaf를 더한 개수도 포함해야 하므로, log(n)의 높이로 구해야함(소수점이 있어 더 필요하니 ceil 사용)
        // Java의 log는 자연로그 함수이므로 log(2)를 통해 나누어 log2로 나누는 효과를 구함
        int h = (int) Math.ceil(Math.log(arrSize) / Math.log(2));

        // Tree에 들어가는 Node의 개수는 2의 높이+1 제곱의 미만 개이다.
        this.treeSize = (int) Math.pow(2, h+1);
        tree = new long[treeSize];
    }

    // 최초 Segment Tree를 구성하는 메소드
    public long init(int[] nums, int node_idx, int start, int end){
        // 만약 start == end라면 leaf 노드라는 의미
        // 즉, 배열의 값을 그대로 저장함
        if(start == end){
            return tree[node_idx] = nums[start];
        }

        // 위에서 return 못했다면 start != end
        // 그렇다면 좌측 노드와 우측 노드의 합으로 구해짐.
        // 좌측 노드의 idx는 *2 이며 우측은 *2+1 이 된다.
        return tree[node_idx] = init(nums, node_idx*2, start, (start+end)/2) +
                init(nums, node_idx*2+1, (start+end)/2+1 ,end);
    }

    // Segment Tree 내 값이 변경되는 경우
    public void update(int node_idx, int start, int end, int idx, long diff){
        // 만약 변경할 index 값이 범위 밖이면 해당 tree는 확인 불 필요
        if(idx < start || end < idx) return;

        // 변경된 차이만큼을 영향 받는 각 node에 더해줌
        tree[node_idx] += diff;

        // leaf 노드에 다다라야 값을 바꿀 수 있으므로 지속 진행
        if(start != end){
            update(node_idx*2, start, (start+end)/2, idx, diff); // 좌측 node로
            update(node_idx*2+1, (start+end)/2+1, end, idx, diff); // 우측 node로
        }
    }

    // 구하고자 하는 덧셈을 진행
    public long sum(int node_idx, int start, int end, int left, int right){
        // 범위를 벗어나게 되는 경우 더할 필요 없음
        if(left > end || right < start){
            return 0;
        }

        // 범위 내 완전히 포함 시에는 더 내려가지 않고 바로 리턴
        if(left <= start && end <= right){
            return tree[node_idx];
        }

        // 그 외의 경우 좌 / 우측으로 지속 탐색 수행
        return sum(node_idx*2, start, (start+end)/2, left, right)+
                sum(node_idx*2+1, (start+end)/2+1, end, left, right);
    }
}

 

세그먼트 트리2 - 느린 전파는 아래 내용을 참고

hongjw1938.tistory.com/25?category=884192

자료구조 - 세그먼트 트리(Segment Tree)2 - 느린 전파

 

이 포스팅은 백준 알고리즘 블로그를 통해 학습하고 문제를 푼 뒤 작성하였습니다. 관련 문제는 아래와 같으며 풀이는 별도로 추후 작성합니다.

백준 알고리즘 2042번: www.acmicpc.net/problem/2042

백준 알고리즘 11505번 : www.acmicpc.net/problem/11505

 

 

오류가 있는 부분은 댓글로 남겨주시면 반영하겠습니다. 감사합니다.