세그먼트 트리(Segment Tree)

세그먼트 트리의 필요성

5칸짜리 배열이 있다고 해보자.
0번째부터 2번째 원소까지의 합은 $1+2+3 = 6$이다.
2번째부터 3번째 원소까지의 합은 $3+4=5$이다.

구간이 주어지면 차례대로 더하면 된다($O(n)$ 소요).
이러한 연산을 $M$번 반복하면 $O(MN)$이 요구된다.

더 빠르게 계산할 수는 없을까?

누적합 방식을 적용해볼 수 있다.

1번째 원소 자리에는 0~1번째 구간의 합,
2번째 원소 자리에는 0~2번째 구간의 합을 넣는 식으로 모든 원소를 채워준다.
=> $[1, 3, 6, 10, 15]$

배열이 이렇게 채워지면,
여기서 0번째부터 2번째 원소까지의 합은 $Array[2]$,
2번째부터 3번째 원소까지의 합은 $Array[3] - Array[1]$이다 => $O(1)$

매우 빠르게 값을 구할 수 있게 되었다.

하지만 여기서 숫자가 주기적으로 변경되는 조건이 붙고,
그 변경된 배열에서의 구간합을 구해야 한다면?

0번째 숫자를 1이 아니라 20으로 변경해보자.
그럼 0번째 숫자뿐 아니라 모든 원소의 값을 변경해야 한다.

이렇게 되면 다시 $O(MN)$이 요구되고, 원점이다.

이럴 때 시간복잡도를 세그먼트 트리를 통해 줄여볼 수 있다.

설명

더 간단하게, $array = [1, 2, 3, 4]$라고 해보자.

각 원소들을 트리에 담을 건데,
우선 리프 노드에 각 원소의 값을 담아준다.

그리고 부모 노드에는 자식 노드들의 값의 합을 담는다.

그러면 아래와 같은 트리가 만들어진다.

(0번째 인덱스는 편의상 생략)

여기서 [1~4]과 같은 표현은
1번째에서 4번째 원소까지의 합을 해당 노드에 담았다는 의미이다.
(이 숫자들까지 트리에 저장할 필요는 없다.)

리프 노드가 4개이므로 높이는 $log4 = 2$이며
필요한 노드의 개수는 $2^3-1$이므로 트리의 크기를 $2^3-1$만큼 만들어야한다.

그림을 예로 들면 2~3 구간합을 구하려고 한다면
2, 3에서 값을 가져오면 된다.
=> $(O(logN))$

트리를 초기화하는 두 가지 방법을 작성해보자면
1. 포화 이진 트리를 만들고, 원소는 전부 최대 높이의 같은 레벨에 둔다
2. 원소는 리프 노드에 두고, 최대 높이가 아니여도 상관 없게 만든다

로직에는 큰 차이는 없으나,
문제에 따라 풀기 수월한 방법을 적용하면 될 것 같다.

구간합을 계산할 때 범위를 어떻게 지정하는지 살펴보자.

발그림을 그려보았다.

targetStart ~ targetEnd 구간의 합을 구한다고 가정하자.
위처럼 5가지 경우가 있겠지만(모서리 부분만 겹치는 경우 등은 그리지 않았음),
구현할 때 5가지 조건 모두 고려해서 작성하지는 않는다.

그림의 의미를 살펴보자. box 형태로 그려놓아서 헷갈릴 수 있는데
그 구간을 나타낸 것 뿐이지 그 이상의 의미는 없다.

그림의 첫 번째 의미를 설명해보면
노드를 탐색하면서 지금 탐색중인 구간이
최종적으로 구해야 할 구간의 합 범위에 아예 들어가있지 않음을 의미한다.

5~8 구간의 합을 구해야 하는데
1~3 구간의 합을 탐색하고 있는 것이다.
아예 겹치는 부분이 없다.

그림의 마지막 부분도 마찬가지이고,

탐색하려는 범위가 아예 목표 범위에 포함되어 있거나,
어느 정도 걸쳐 있거나 등의 경우가 있는 것이다.

범위가 아예 포함되어 있는 경우만 노드의 값을 반환하고,
범위가 걸쳐 있는 경우에는 한 단계 더 탐색 해보는 것이다.

트리 초기화 코드

리프노드의 첫 부분의 위치를 계산해
미리 원소를 트리에 입력해 놓고 함수를 실행한다.

리프노드가 8번째부터 시작한다고 하면
트리의 8번째부터 원소를 입력 받는 것이다.

이렇게 리프노드가 시작되는 위치인 S를 계산해야 한다.

하지만 2번 방법을 이용하면
리프 노드의 시작 지점을 계산할 필요 없다.

(트리를 초기화 하면 리프노드로 원소가 이동하는 것은 동일하고,
원소 값을 임시 배열에 저장해놓고 트리에 옮기는데,
트리의 리프노드 원소 값을 변경하면 임시 배열 원소 값도 변경하자!)

// 1번 방법

void createTree(int current, int start, int end) {
    if(start == end)
        return arr[current];
    
    int mid = (start + end) / 2;
    
    arr[current] = createTree(current * 2, start, mid) 
                 + createTree(current * 2, mid + 1, end);
        
}

//2번 방법

static long createTree(int current, int start, int end) {
        if (start == end)
            return tree[current] = arr[start];

        int mid = (start + end) / 2;

        return tree[current] = createTree(current * 2, start, mid) + createTree(current * 2 + 1, mid + 1, end);
    }

트리 값 수정 코드

포화 이진 트리 상태에서 반복문을 이용해 값을 update 하는 것도 매우 간단하다.
(리프 노드에서 값을 수정하고, 나누기 2를 반복해가며 조상들의 값을 수정한다.)

여기서 diff는 값의 차이를 의미한다.
예를 들어 값 a를 값 b로 변경해야 하면,
b를 인자로 넘기는 것이 아니라
b - a를 넘긴다. 값의 차를 넘겨야 계산이 편해진다.
((b - a) + a는 b니까!)

//1번 방법
static void updateTree(int current, int start, int end, int index, long diff) {

        int mid = (start + end) / 2;

        if(current == index) {
            arr[current] += diff;
            return;
        }

        if(start <= changeIndex && changeIndex <= end) {
            arr[current] += diff;
            updateTree(current * 2, start, mid, index, diff);
            updateTree(current * 2 + 1, mid + 1, end, index, diff);
        }
    }
    
 //2번 방법
 static void updateTree(int current, int start, int end, int changeIndex, long diff) {

        if (changeIndex < start || changeIndex > end)
            return;

        tree[current] += diff;
        int mid = (start + end) / 2;

        if (start != end) {
            updateTree(current * 2, start, mid, changeIndex, diff);
            updateTree(current * 2 + 1, mid + 1, end, changeIndex, diff);
        }
    }
   

구간합 계산 코드

static long getSum(int current, int currentStart, int currentEnd, int targetStart, int targetEnd) {

        int mid = (currentStart + currentEnd) / 2;

        if (currentEnd < targetStart || currentStart > targetEnd)
            return 0;

        if (currentStart >= targetStart && currentEnd <= targetEnd)
            return tree[current];

        return getSum(current * 2, currentStart, mid, targetStart, targetEnd)
                + getSum(current * 2 + 1, mid + 1, currentEnd, targetStart, targetEnd);

    }

관련 문제

기본 문제

구간 합 구하기
단순히 세그먼트 트리에 대한 개념을 안다면 쉽게 풀 수 있다.

포화이진트리로 만들지 않고 풀었는데

if (a == 1) {
                updateTree(1, 1, n, b, c - arr[b]);
                arr[b] = c; // 이 부분 안쓴걸 파악 못해서 시간 엄청 날렸다..
            }

값을 수정했을 때 초기에 입력받은 배열의 값도 변경했어야 했는데
그렇지 못해서 계속 AC를 받지 못했다.

침착하게 풀어야 한다.

커피숍2
이것도 세그먼트 트리를 구현하면 바로 풀리는 문제이다.

근데 또 한가지 실수를 했다.

위의 구간 합 구하기 문제에서와 똑같은 부분에서 실수를 했는데,

if (a == 1) {
		arr[b] = c;
                updateTree(1, 1, n, b, c - arr[b]);
                 
            }

순간적으로 이렇게 작성해도 똑같을거라 생각했다.

원래는 이렇게 써야 한다.

if (a == 1) {
		// arr[b] = c; 여기가 아니라
                updateTree(1, 1, n, b, c - arr[b]);
                arr[b] = c; // 여기에 작성
            }

당연히 다른 결과를 줄텐데 순간적으로 착각했다.
arr[b]에 c를 대입하고 c - arr[b]를 인자로 넘기면
무조건 값 0을 넘겨줄 수 밖에 없기 때문이다.

쿼리를 발생시키고 처음에 입력받은 배열의 값을 변경하자.

최솟값과 최댓값
이것도 쉽다.. 위의 문제들보다 더 쉽다.

최댓값과 최솟값을 한 번에 저장하는 트리를 만들거나
최댓값, 최솟값을 따로 저장하는 트리를 만들면 된다(2개의 트리 생성).

응용 문제

사탕상자
살짝 응용된 문제이다.

맛이 $10^6$가지이고 총 사탕 개수는 20억을 넘지 않는다.
각 노드에 사탕의 총 개수를 저장해주면

한 노드의 값이 9이고,
왼쪽 노드 값이 4, 오른쪽 노드가 5일 때

6번째 순위의 사탕을 꺼내려고 하면
오른쪽 노드로 이동해야 한다.

왼쪽 노드들에 해당하는 사탕 중에는 찾으려는 사탕이 없고,
오른쪽 노드들에 해당하는 사탕 중에서 찾으려는 사탕이 있다는 뜻이다.

이때 순위의 값 6을 그대로 들고 가는 것이 아니라,
왼쪽 노드 값인 4를 뺀 값 2를 들고 가야 한다.

값을 왜 뺄까?

예를 들어보자.
전교생이 9명일 때,
A반에 1~4등이 소속되어 있고
B반에 5~9등이 소속되어 있다고 하자.

여기서 6등인 학생 S한테 찾아가려고 하면,
당연히 B반에 찾아갈 것인데

학생 S는 B반 안에서도 6등일까?
B반 안에서 학생 S는 2등이다. A반 학생들은 제외해야 한다.

값을 빼지 않는다는 것은 S가 B반에서도 6등이라고 보는 것과 같은 의미이다.