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

세그먼트 트리(Segment Tree)는 부분 합을 트리 구조에 저장해서, 값이 바뀔 수 있는 상황에서도 부분합을 빠르게 구할 수 있는 방법이다.

 

세그먼트 트리를 그려보면 이런 느낌이다.

위의 대괄호 안에 들어있는 숫자는 각 노드의 인덱스이고, 아래 큰 숫자는 부분합의 범위이다.

즉 루트 노드인 [0]에는 배열의 0~9번째 값들의 합이 저장되어 있고, [11] 노드에는 배열의 5~6번째 값들의 합이 저장되어 있는 식이다.

---

구현 |

 

1. init

배열의 초기 값을 세그먼트 트리에 처음 넣는 함수이다.

이때는 루트 노드 [0]부터 아래로 내려간다.

한 노드에 값을 채울 때는 자식 노드에 값을 먼저 채우고 그 뒤에 그 합을 쓰면 되는 것이다.

그런 식으로 말단 노드 (예를 들어, 6번째부터 6번째까지의 합 노드 = 그냥 6번째 값 노드) 까지 내려갔다 오면 된다.

루트 노드가 [0]일 때 [index]의 자식 노드는 [index * 2 + 1], [index * 2 + 2]다.

이를 이용해서 일단 이렇게 간단하게 짜면 된다.

int init(int S, int E, int index)
{
    if (S == E)
    {
        return segTree[index] = array[S];
    }
    return segTree[index] = init(S, (S + E) / 2, index * 2 + 1) + init((S + E) / 2 + 1, E, index * 2 + 2);
}

어떤 배열의 0번째부터 n번째까지의 값을 세그트리에 넣고 싶다면 init(0, n, 0); 을 쓰면 된다.

 

2. change

배열 내에서 값이 변하면 트리에도 적용시켜주는 함수가 필요하다.

이 때는 반대로 말단 노드에서 위로 계속 올라가야 하는데, 내가 바꾸려는 값이 들어있는 말단 노드가 어디인지 못 찾겠다.

그래서 init할 때 해당 배열의 값이 트리의 어느 인덱스에 있는지 저장하는 변수를 하나 만들어준다.

int init(int S, int E, int index)
{
    if (S == E)
    {
        treeIndex[S] = index;
        return segTree[index] = array[S];
    }
    return segTree[index] = init(S, (S + E) / 2, index * 2 + 1) + init((S + E) / 2 + 1, E, index * 2 + 2);
}

한 줄 추가해주면 된다.

그럼 위의 treeIndex를 갖고 말단 노드 값을 바꿔주고, 그 부모 노드도 바꿔주고, 다시 그 부모 노드도 바꿔줄 수 있다.

void change(int index, int newVar)
{
    array[index] = newVar;
    segTree[treeIndex[index]] = newVar;
    int i = (treeIndex[index] - 1) / 2;
    while (i >= 0)
    {
        segTree[i] = segTree[i * 2 + 1] + segTree[i * 2 + 2];
        if (i == 0)
            break;
        else
            i = (i - 1) / 2;
    }
}

index는 바꿔줄 배열의 값의 index고, newVar은 새로운 값이다.

우선 array와 말단 노드 값을 바꿔준 후 루트 노드까지 계속 올라가며 바꿔주면 된다.

 

3. sum

결론적으로 구간 합을 구해야 하는데, 예를 들어 2~7번째 배열의 부분합을 구해야 한다면 세그트리에서는 아래 표시해놓은 노드들의 값을 더하면 되는 접근이다.

 

int sum(int start, int end, int S, int E, int index)
{
    if (start > E || end < S)
    {
        return 0;
    }
    else if (start >= S && end <= E)
    {
        return segTree[index];
    }
    else
    {
        return sum(start, (start + end) / 2, S, E, index * 2 + 1) + sum((start + end) / 2 + 1, end, S, E, index * 2 + 2);
    }
}

start, end는 현재 함수에서 보고 있는 구간이고, S와 E는 합을 구하고 싶은 구간이다. 그래서 함수를 호출할 때도 S와 E는 변하지 않고 start와 end만 바꿔준다.

지금 보고 있는 구간이 S~E에 포함되어 있으면 그냥 그 구간의 합을 반환하고,

반쪽만 걸쳐(?) 있으면 지금 구간을 반으로 나눠서 다시 함수에 넣어준다.

 

응용하면 구간 합 뿐만 아니라 최소/최댓값 등 많은 부분에서 유용하게 쓸 수 있다. 세그트리 조아~

---

관련 문제 |

[BOJ] 구간 합 구하기 (https://www.acmicpc.net/problem/2042)

[BOJ] 최솟값 (https://www.acmicpc.net/problem/10868)

[BOJ] 최솟값과 최댓값 (https://www.acmicpc.net/problem/2357)