[Algorithm] 세그먼트 트리

tngus2sh·2024년 9월 1일
alogrithm

Algorithm

목록 보기
13/18

세그먼트 트리를 배우기 전..

/icons/compose_gray.svg 크기가 N인 정수 배열 A가 있고, 여기서 다음과 같은 연산을 최대 M번 수행해야 하는 문제가 있습니다.
  1. 구간 l, r(l ≤ r)이 주어졌을 때, A[l] + A[l + 1] + … + A[r - 1] + A[r]을 구해서 출력하기
  2. i번째 수를 v로 바꾸기 (A[i] = v)

1번 연산 A[l] + A[l + 1] + … + A[r - 1] + A[r]을 구하기 위해 소스 1과 같이 모두 더하는 방법이 있습니다.

int ans = 0;
for (int i = l; i <= r; i++) {
	ans += a[i];
}

해당 코드의 시간 복잡도는 O(N)입니다. 2번 연산 A[i] = v는 O(1)입니다. 연산을 최대 M번 수행해야 하니 연산 하나의 시간 복잡도는 O(N)입니다. 총 시간 복잡도는 O(NM)입니다.

누적 합을 사용하면, 1번 연산의 시간 복잡도를 O(1)로 줄일 수 있습니다. 하지만, 2번 연산으로 수가 변경될 때마다 누적 합을 다시 구해야 하기 때문에, 2번 연산의 시간 복잡도는 O(N)입니다. 연산 하나의 시간 복잡도는 O(N)이니 총 시간 복잡도는 O(NM)이 됩니다.

세그먼트 트리를 사용하면 해당 연산들의 시간 복잡도를 줄일 수 있습니다.

세그먼트 트리(Segment Tree)

여러 개의 데이터가 존재할 때 특정 구간의 합(최솟값, 최댓값, 곱 등)을 구하는 데 사용하는 자료구조
트리 종류 중에 하나로 이진 트리의 형태이며, 특정 구간의 합을 가장 빠르게 구할 수 있다.

세그먼트 트리를 사용하면 1번 연산과 2번 연산을 O(logN)에 수행할 수 있다.

  • 리프 노드 : 배열의 그 수 자체
  • 리프 노드가 아닌 노드 : 왼쪽 자식과 오른쪽 자식의 합을 저장

어떤 노드의 번호가 x일 때, 왼쪽 자식의 번호는 2x, 오른쪽 자식의 번호는 2x + 1이 된다.

과정

  • 리프 노드를 제외한 다른 모든 노드는 항상 2개의 자식을 가진다.
    • 따라서, 세그먼트 트리는 Full Binary Tree의 형태를 가진다.
    • 만약, N이 2의 제곱꼴인 경우에는 Perfect Binary Tree가 된다.
  • 리프 노드가 N개인 Full Binary Tree에는 리프노드가 아닌 노드가 N-1개 있다. 따라서, 필요한 노드의 수는 2N-1개이다. /icons/light-bulb_gray.svg 참고로 필요한 노드의 수랑 배열의 크기랑은 별개이다. → 세그먼트 트리의 경우 Full Binary Tree인데, 완전 이진 트리가 아니기 때문에 배열의 순서대로 노드가 채워지는 것이 아니다. - N이 2의 제곱꼴이 아닌 경우에 높이 h = logN이다.
  • 세그먼트 트리의 정보를 저장하기 위해서 배열을 사용한다. 깊이가 가장 깊은 노드와 가장 깊지 않은 리프 노드의 깊이 차는 1보다 작거나 같다. 따라서, 배열을 이용해도 공간을 크게 낭비하지 않는다.
    • tree[x]에 노드 x의 정보를 저장
/icons/light-bulb_gray.svg **Size 구하는 방법** 먼저 h를 구하고 tree_size를 결정하면 된다. 
int h = (int)Math.ceil(Math.log(n) / Math.log(2));
int tree_size = 1 << (h + 1);
// a: 배열 A
// tree: 세그먼트 트리
// node: 노드 번호
// node에 저장되어 있는 합의 범위가 start - end
void init(long[] a, long[] tree, int node, int start, int end) {
		// 리프 노드
    if (start == end) {
        tree[node] = a[start];
    } else {
        init(a, tree, node*2, start, (start+end)/2);
        init(a, tree, node*2+1, (start+end)/2+1, end);
        tree[node] = tree[node*2] + tree[node*2+1];
    }
}

start == end인 경우는 리프노드의 경우이다. 리프 노드는 배열의 그 수를 저장해야 하기 때문에, tree[node] = a[start]가 된다.

node의 왼쪽 자식은 node2이고, 오른쪽 자식은 node2 + 1이다. 또, node에 저장된 구간이 [start, end]라면, 왼쪽 자식은 [start, (start + end)/2], 오른쪽 자식은 [(start + end)/2 + 1, end]가 저장된 구간이다.

tree[node]에 저장될 값을 구하려면 왼쪽 자식에 저장된 값 tree[node*2], 오른쪽 자식에 저장된 값 tree[node*2 + 1]을 먼저 구해야한다. 따라서, 재귀 함수를 이용해 각각의 값을 먼저 구한다.

구간의 합 구하기

구간 left, right가 주어졌을 때, 합을 구하려면 트리를 순회하면서 각 노드에 저장된 구간의 정보와 left, right와의 관계를 살펴봐야 한다.

node에 저장된 구간이 [start, end]이고, 합을 구해야 하는 구간이 [left, right] 라면 다음과 같이 4가지 경우로 나누어질 수 있다.

  • [left, right]와 [start, end]가 겹치지 않는 경우 - ①
  • [left, right]가 [start, end]를 완전히 포함하는 경우 - ②
  • [start, end]가 [left, right]를 완전히 포함하는 경우 - ③
  • [left, right]와 [start, end]가 겹쳐져 있는 경우(1,2,3 제외한 나머지 경우) - ④

① if (left > end || right < start)로 나타낼 수 있다.

  • left > end : [start, end] 뒤에 [left, right]가 있는 경우
  • right < start : [start, end]앞에 [left, right]가 있는 경우

이 경우에는 겹치지 않기 때문에 더 이상 탐색을 이어나갈 필요가 없다. 따라서 0을 리턴해 탐색을 종료한다.

② if(left ≤ start && end ≤ right)로 나타낼 수 있다.

이 경우도 더 이상 탐색을 이어나갈 필요가 없다. 구해야 하는 합의 범위는 [left, right]인데, [start, end]는 그 범위에 모두 포함되고, 그 node의 자식도 모두 포함되기 때문에 더이상 호출을 하는 것은 비효율적이다. 따라서, tree[node]를 리턴해 탐색을 종료한다.

③, ④ 왼쪽 자식과 오른쪽 자식을 루트로 하는 트리에서 다시 탐색을 시작해야 한다.

Untitled

예를 들어 위와 같은 경우를 보자.

left = 0, right = 9인 경우에는 맨 위의 루트 노드만으로 바로 합을 알 수 있다.

하지만, left = 2, right = 4인 경우 다음과 같이 구해진다.

  1. 루트 노드에서 재귀 함수의 형태로 탐색을 시작 왼쪽 오른쪽 탐색이 시작된다. 여기서 오른쪽 탐색은 ①번의 경우로서 겹치는 부분이 없어서 탐색을 종료한다.
  2. 0-9에서 왼쪽으로 탐색이 들어가고 0-4로 갈 것이다. 구하려는 합의 범위가 2-4이기 때문에, 0-4가 더 큰 범위라서 나눠져서 들어가야 한다.(③의 경우)
  3. 그래서 왼쪽 오른쪽으로 다시 탐색이 들어간다. 이와 같은 방식으로 2, 3-4의 합을 찾아내서 그 둘을 더해 구간의 합을 구한다. 여기서 5-9는 ①의 경우로서 겹치지 않기 때문에 탐색을 종료한다.

left = 3, right = 9인 경우에는 위와 같은 방식으로 루트 노드에서 왼쪽 서브트리와 오른쪽 서브 트리로 탐색이 이루어진다. 여기서 주목해야 할 점은 오른쪽 서브트리이다.

  1. 5-9의 경우에는 3-9에 포함되어 있기 때문에 결국 더이상의 탐색이 불필요하다. 이미 구하려는 범위의 합에 포함되어 있기 때문(②의 경우)
  2. 최종적으로 왼쪽 서브 트리에서 3-4, 오른쪽 서브트리에서 5-9를 꺼내와서 둘의 합을 구하면서 종료한다.
// 배열의 특정 구간 합을 세그먼트 트리로 구하기
long sum(int node, int start, int end, int left, int right) {
    // 노드가 가지는 값의 구간이 구하려고 하는 합의 구간에 속하지 않는 경우 0리턴
    if (end < left || right < start) {
        return 0;
    } else if (left <= start && end <= right) {
        // 노드가 가지는 값의 구간이 구하려고 하는 합의 구간에 속하는 경우 노드 값 리턴
        return tree[node];
    } else {
        // 그 외는 2가지 경우가 존재
        // 1. 노드가 가지는 값의 구간이 구하려고 하는 합의 구간에 일부는 속하고 일부는 속하지 않는 경우
        // 2. 노드가 가지는 값의 구간이 구하려고 하는 합의 구간을 모두 포함하는 경우
        // 이와 같은 경우에는 자식노드를 탐색해서 값을 리턴
        return sum(node*2, start, (start+end)/2, left, right)
                + sum(node*2+1, (start+end)/2+1, end, left, right);
    }
}

수 변경하기

index번째 수를 val로 변경하는 경우, index번째를 포함하는 노드에 들어있는 합만 변경해주면 된다.

원래 index번째 수가 a[index]였고, 바뀐 수가 val이라면 합은 val - a[index]만큼 변한다.

수 변경은 2가지 경우가 있다.

  1. [start, end]에 index가 포함되는 경우
  2. [start, end]에 index가 포함되지 않는 경우

1번 경우에는 재귀호출을 진행하고, 2번의 경우는 그 노드의 모든 자식도 index번째를 포함하지 않으니 재귀 호출을 중단하면 된다.

리프 노드를 찾으면 그 노드의 합을 변경해준다. 이후 리턴될 때마다 각 노드의 합을 자식에 저장된 합을 이용해 다시 구하면 된다.

void update(int node, int start, int end, int index, long diff) {
    if (index < start || end < index) {
        return;
    }else {
        // 노드가 가지는 값의 구간에 배열의 인덱스(값이 변경 될 인덱스)가 포함되는 경우
        // 노드의 값 + 차이값(변경할 값-기존값)
        tree[node] = tree[node] + diff;

        // 노드가 리프노드가 아닌 경우
        if (start != end) {
            // 리프노드까지 계속 자식노드를 탐색
            update(node*2, start, (start+end)/2, index, diff) ;
            update(node*2+1, (start+end)/2+1, end, index, diff) ;
        }
    }
}

구현

  • 전체 코드
    static class SegmentTree{
        // 세그먼트 트리를 구현할 배열
        private long[] tree;

        // 생성자에서 세그먼트 트리의 전체노드 수 계산 (즉, 배열 길이)
        SegmentTree(int n) {
            // 트리의 높이 계산
            double treeHeight = Math.ceil(Math.log(n)/Math.log(2))+1;
            // 트리의 노드 수 계산
            long treeNodeCount = Math.round(Math.pow(2, treeHeight));
            // 트리의 길이 설정
            tree = new long[Math.toIntExact(treeNodeCount)];
        }

        // 세그먼트 트리의 노드 값 초기화
        long init(long[] arr, int node, int start, int end ){
            // 세그먼트 트리의 리프노드인 경우
            if (start == end) {
                // 리프노드에 배열의 값 저장 후 리턴
                return tree[node] = arr[start];
            }else{
                // 리프노드가 아닌 경우에는 자식노드의 값을 더해서 노드의 값 초기화 후 리턴
                return tree[node] = init(arr, node*2, start, (start+end)/2)
                        + init(arr, node*2+1, (start+end)/2+1, end);
            }
        }

        // 배열의 특정 구간 합을 세그먼트 트리로 구하기
        long sum(int node, int start, int end, int left, int right) {
            // 노드가 가지는 값의 구간이 구하려고 하는 합의 구간에 속하지 않는 경우 0리턴
            if (end < left || right < start) {
                return 0;
            } else if (left <= start && end <= right) {
                // 노드가 가지는 값의 구간이 구하려고 하는 합의 구간에 속하는 경우 노드 값 리턴
                return tree[node];
            } else {
                // 그 외는 2가지 경우가 존재
                // 1. 노드가 가지는 값의 구간이 구하려고 하는 합의 구간에 일부는 속하고 일부는 속하지 않는 경우
                // 2. 노드가 가지는 값의 구간이 구하려고 하는 합의 구간을 모두 포함하는 경우
                // 이와 같은 경우에는 자식노드를 탐색해서 값을 리턴
                return sum(node*2, start, (start+end)/2, left, right)
                        + sum(node*2+1, (start+end)/2+1, end, left, right);
            }
        }

        // 배열의 특정 인데스의 값이 변경 될 경우 세그먼트 트리의 노드 값 변경(차이 값을 더하는 방법)
        void update(int node, int start, int end, int index, long diff) {
            // 노드가 가지는 값의 구간에 배열의 인덱스(값이 변경 될 인덱스)가 포함되지 않을 경우7
            if (index < start || end < index) {
                // 아무것도 안함
                return;
            }else {
                // 노드가 가지는 값의 구간에 배열의 인덱스(값이 변경 될 인덱스)가 포함되는 경우
                // 노드의 값 + 차이값(변경할 값-기존값)
                tree[node] = tree[node] + diff;

                // 노드가 리프노드가 아닌 경우
                if (start != end) {
                    // 리프노드까지 계속 자식노드를 탐색
                    update(node*2, start, (start+end)/2, index, diff) ;
                    update(node*2+1, (start+end)/2+1, end, index, diff) ;
                }
            }
        }

        // 배열의 특정 인데스의 값이 변경 될 경우 세그먼트 트리의 노드 값 변경(노드 값을 직접 변경)
        long update2(int node, int start, int end, int index, long changeValue) {
            // 노드가 가지는 값의 구간에 배열의 인덱스(값이 변경 될 인덱스)가 포함되지 않을 경우
            if (index < start || end < index) {
                // 트리의 노드 값 리턴
                return tree[node];
            } else if (start == index && end == index) {
                // 노드가 가지는 값의 구간과 배열의 인덱스(값이 변경 될 인덱스)값이 같은 경우
                // 노드의 값을 변경 될 값으로 변경
                return tree[node] = changeValue;
            } else {
                // 노드가 가지는 값의 구간에 배열의 인덱스(값이 변경 될 인덱스)값이 포함되는 경우(같은 경우는 제외)
                // 자식 노드를 탐색 후 값을 더해서 리턴
                return tree[node] = update2(node * 2, start, (start + end) / 2, index, changeValue) +
                        update2(node * 2 + 1, (start + end) / 2 + 1, end, index, changeValue);
            }
        }
}

시간 복잡도

구간의 합

트리의 각 레벨에서 방문하는 노드의 개수는 4개를 넘지 않는다. 트리의 높이 h는 logN이기 때문에, 합을 구하는 시간 복잡도는 logN이다.

첫번째 레벨에서는 루트 노드 하나만 있고, 루트 노드는 반드시 방문하게 된다.

리프 노드가 아닌 노드는 2개의 자식을 갖고, 재귀 호출을 하는 경우 항상 2개의 호출을 하게 된다. 어떤 레벨에서 방문한 노드의 개수가 2개 이하인 경우에는 다음 레벨에서 방문한 노드의 개수는 4개 이하이다.

어떤 레벨에서 방문한 노드의 수가 3개 또는 4개인 경우에 다음 레벨에서 방문한 노드의 수가 4개 이하인지 살펴보면 된다.

다음 페이지에서 소스 3에서 예제를 입력해보면서 동작하는 방식을 확인해보는 것이 있으니 참고

세그먼트 트리

레벨 l에서 방문한 노드가 3개이고 l,m,r이라고 하자. m은 절대로 재귀 호출을 하지 않는다. 세그먼트 트리의 모든 쿼리는 연속된 구간의 합을 구하게 되는데, m은 항상 부모 노드의 구간에 포함되는 노드이다. 재귀 호출이 일어났다는 것은 그 구간의 일부만 포함되어야 한다는 것을 의미하기 때문. 방문한 노드가 4개인 경우도 가장 왼쪽과 오른쪽에 있는 재귀 호출을 할 수 있고, 가운데 있는 노드는 재귀 호출을 하지 않는다.

따라서, 각 레벨에서 최대 4개의 노드만 방문할 수 있다.

T(n)=O(logN)T(n) = O(logN)

수 변경하기

트리의 각 레벨에서 방문하는 노드의 개수는 2개를 넘지 않는다. 트리의 높이 H는 logN이기 때문에, 시간 복잡도는 logN이다.

T(n)=O(logN)T(n) = O(logN)

세그먼트 트리를 이용해서 구간의 최솟값도 구할 수 있다. 최솟값 이외에도 최댓값, 곱, XOR 연산 등도 구할 수 있다.

profile
tngus2sh
백엔드 개발자
이전 포스트

[Algorithm] Prim

다음 포스트

[Algorithm] Binary Tree

0개의 댓글