[WEEK04] 세그먼트 트리

세그먼트 트리란?

세그먼트 트리는 배열의 구간에 대한 정보를 저장하는 트리 자료구조이다. 주로 다음과 같은 경우에 사용된다:

• 부분 구간의 합을 자주 구해야 할 때
• 배열의 일부 값이 자주 갱신될 때
• 최대/최솟값을 빠르게 구해야 할 때

이전에는 단순히 리스트 슬라이싱과 sum() 함수를 썼는데, 쿼리와 업데이트가 많아질수록 비효율적이라는 걸 알게 되었다. 이때 세그먼트 트리를 사용하면 다음과 같은 시간 복잡도를 얻을 수 있다.

• 구간 쿼리: O(log N)
• 값 업데이트: O(log N)

기본 아이디어

• 세그먼트 트리는 완전 이진 트리 형태로 구성된다.
• 리프 노드는 원래 배열의 값, 내부 노드는 해당 구간의 정보를 담는다.
• 예를 들어 구간 합 세그먼트 트리라면, 내부 노드는 자식 노드들의 합을 저장한다.

💻 코드 구현

n = int(input())
k = 0
length = n

while length != 0:
    length //= 2
    k += 1

size = 2**k * 2
start_idx = 2**k - 1
tree = [0] * (size+1)

for i in range(start_index+1, start_index+n+1):
    tree[i] = int(input())

def create_tree(i):
    while i != 1:
        tree[i//2] += tree[i]
        i -= 1

create_tree(size-1)

def get_sum(s, e):
    result = 0
    while s <= e:
        if s % 2 == 1:
            result += tree[s]
            s += 1
        if e % 2 == 0:
            result += tree[e]
            e -= 1
        s //= 2
        e //= 2
    return result

def modify(idx, val):
    dv = val - tree[idx]
    while idx > 0:
        tree[idx] += dv
        idx //= 2

정리

• 세그먼트 트리는 구간 쿼리와 업데이트에 효율적이다.
• 구현이 다소 복잡하지만 한 번 잘 만들어 두면 여러 문제에 재활용 가능하다.
• 오늘은 구간 합만 다뤘지만, 최소값, 최대값, 곱셈 등 다른 연산도 비슷한 방식으로 구현할 수 있다.

📚 참고자료

• [CP Algorithms] Segment Trees
• 알고리즘 코딩테스트 핵심이론 강의 - 세그먼트트리(인덱스 트리)