세그먼트 트리

소개

세그먼트 트리는 구간에 대한 연산을 $O(logN)$만에 처리할 수 있는 자료구조입니다. 난이도 높은 문제에서 다양한 방식으로 활용되는 자료구조이기도 합니다. 

 

구조

이진 트리의 형태이며, 각 노드는 특정 구간의 값을 저장하고 있습니다. 이 값은 구간의 합, 최솟값 등 다양한 값이 될 수 있습니다. 루트 노드는 전체 구간에 대한 값을 저장하며, 루트 노드에서 멀어질수록 담당하는 구간이 작아지다 리프 노드에 도달하면 원소 하나만을 담당하게 됩니다. 

전체 구간이 1 ~ 7일 때 세그먼트 트리의 각 노드가 담당하는 구간

 

연산

기본적으로 원소 하나를 변경하는 연산과 특정 구간의 값을 구하는 연산을 처리할 수 있습니다. 

구간 합을 구하는 세그먼트 트리의 코드를 예시로 설명하겠습니다. 

원소 변경

void change(int now, int l, int r, int idx, long long diff) {
    if (idx < l or r < idx)
        return;
    tree[now] += diff;
    if (l == r) 
        return;
    int mid = (l + r) / 2;
    change(now * 2, l, mid, idx, diff);
    change(now * 2 + 1, mid + 1, r, idx, diff);
}

now는 현재 노드의 번호입니다. 세그먼트 트리는 완전 이진 트리에 가까운 이진 트리이며, 따라서 배열에 값을 저장할 수 있습니다. tree 배열으로 트리를 표현하고 있습니다. idx는 변경시킬 값의 인덱스이며, diff은 변경된 값과 이전의 값의 차입니다. l과 r은 현재 노드가 담당하는 구간입니다. 만약 현재 노드의 구간이 변경시킬 값을 포함하지 않는다면 함수를 종료하고, 아니라면 현재 구간의 값을 변경시킨 뒤 자식 노드들의 값을 변경시킵니다. 이때 자식 노드가 담당하는 구간은 각각 [l, (l + r) / 2], [(l + r) / 2 + 1, r]이 됩니다. 

 

구간의 값 구하기

long long getsum(int now, int nowl, int nowr, int l, int r) {
    if (r < nowl or nowr < l)
        return 0;
    if (l <= nowl and nowr <= r)
        return tree[now];
    int mid = (nowl + nowr) / 2;
    return getsum(now * 2, nowl, mid, l, r) + getsum(now * 2 + 1, mid + 1, nowr, l, r);
}

change함수와 비슷한 느낌입니다. 

now는 현재 노드의 번호이며, nowl, nowr은 현재 노드가 담당하는 구간입니다. l, r은 구하고자 하는 구간입니다. 

만약 구하고자 하는 구간이 현재 노드의 구간과 겹치지 않으면 0을, 현재 노드의 구간에 완전히 포함된다면 현재 노드의 값을 바로 반환합니다. 현재 노드의 구간과 구하고자 하는 구간이 일부 겹친다면 자식 노드들을 통해 겹치는 구간의 합을 구하여 반환합니다. 

 

정당성 증명

1. 항상 정답을 구할 수 있을까? 

getsum 함수는 노드의 구간이 구하고자 하는 구간에 완전히 포함될 때까지 재귀호출합니다. 살펴보는 노드가 점점 아래로 내려가다 리프 노드에 도달하면 구하고자 하는 구간과 나머지 구간이 명확히 구분되며, 구하고자 하는 구간의 값만 반환되어 위로 올라가기 때문에 항상 올바른 정답을 구할 수 있습니다. 

 

2. 시간 복잡도가 어떻게 될까?

세그먼트 트리가 구하고자 하는 구간은 연속된 구간의 값입니다. 따라서 같은 높이에 있는 노드들 중  자신의 자식 노드를 살펴보아야 하는 노드는 2개를 넘을 수 없습니다. 2개를 넘는 순간 중간에 위치한 노드의 범위는 구하고자 하는 범위에 완전히 포함되어 즉시 자신의 값을 반환하기 때문입니다. 구간의 원소 수를 N이라고 하면 쿼리의 시간복잡도는 트리의 높이에 비례하는 $O(logN)$이 되며, 이 정도의 시간복잡도면 무리없이 문제를 해결할 수 있습니다. 

 

추가

• 속도는 조금 느려지지만, Lazy propagation을 사용하여 구간 전체의 값을 변경할 수 있습니다. (시간 복잡도는 변하지 않습니다)
• 트리가 저장하는 값과 자식 노드의 결과를 바탕으로 자신의 값을 계산하는 과정을 손보면 다양한 문제에 활용할 수 있는 자료구조입니다.