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

세그먼트 트리(Segment Tree) 자료구조는 누적 합 알고리즘의 심화 버전으로 이해될 수 있다.

 

크기가 N인 배열이 주어졌을 때, 누적 합 알고리즘의 시간복잡도는 O(N)이다. 이도 빠르긴 하지만, N의 크기가 더 커졌을 경우에는 이 또한 시간 초과가 나올 수 있다. 세그먼트 트리라는 완전 이진 트리 형태의 자료구조를 사용하면 배열의 특정 구간의 합을 구하는 시간복잡도는 O(logN)으로 훨씬 더 시간을 단축할 수 있다.

 

 

배열의 크기가 N일 경우의 세그먼트 트리의 형태은 다음과 같다.

하나의 노드마다 인덱스 i~인덱스 j 사이인 배열 구간의 누적 합 결과값이 들어간다.

 

 

 

 

 세그먼트 트리 만들기

완전 이진 트리는 배열을 통해 구현한다.

어떠한 부모 노드의 두 자식 노드들은 부모 노드의 인덱스를 idx라고 했을 때, 각각 tree[idx*2], tree[idx*2+1]의 위치에 저장된다. 따라서, 세그먼트 트리를 구현하기 위해서는 배열의 크기가 N이라고 했을 때, 넉넉하게 N*4 크기의 배열이 필요하다.

이때, 주의해야 할 점은 트리의 루트 노드의 인덱스는 1부터 시작한다.

 

트리를 만드는 과정은 재귀 알고리즘을 사용한다.

 

start: 배열에서 현재 가리키는 왼쪽 인덱스

end: 배열에서 현재 가리키는 오른쪽 인덱스

idx: 트리에서 현재 가리키는 인덱스

 

인덱스 start와 인덱스 end의 위치가 같다면, 가리키는 트리의 인덱스의 값은 배열 start의 값으로 정한다.

아직 위치를 찾지 못했다면 mid를 기준으로 양쪽을 나눈 뒤, 두 자식 노드로 내려가 탐색을 계속해서 진행한다.

 

arr = [0, 1, 2, 3, 4, 5]
tree = [0 for _ in range(len(arr)*4)]


def init(start, end, idx):
    if start == end:
        tree[idx] = arr[start]
        return tree[idx]

    mid = (start + end) // 2
    tree[idx] = init(start, mid, idx*2) + init(mid+1, end, idx*2+1)
    return tree[idx]

 

 

 구간 합 구하기

start: 배열에서 현재 가리키는 왼쪽 인덱스

end: 배열에서 현재 가리키는 오른쪽 인덱스

idx: 트리에서 현재 가리키는 인덱스

left: 구하고자 하는 특정 구간의 왼쪽 인덱스

right: 구하고자 하는 특정 구간의 오른쪽 인덱스

 

인덱스 left~right까지의 구간의 합을 구하기 위한 코드를 구현한다.

구하고자 하는 구간의 범위 밖에 있다면 0을 반환하고, 완전히 내부에 있다면 트리의 인덱스값을 반환한다.

그 외에는 아직 분할이 덜 되었다는 의미이므로 mid인덱스를 기준으로 왼쪽 구간과 오른쪽 구간을 나누어 분할을 계속한다.

 

def interval_sum(start, end, idx, left, right):
    if left > end or right < start:
        return 0
    if left <= start and end <= right:
        return tree[idx]

    mid = (start + end) // 2
    return interval_sum(start, mid, idx*2, left, right) + interval_sum(mid+1, end, idx*2+1, left, right)

 

 

 트리의 요소를 변경할 경우

트리의 요소를 변경하게 된다면, 해당 위치와 관계가 있는 구간들의 값들을 모두 바꾸어주어야 한다.

 

start: 배열에서 현재 가리키는 왼쪽 인덱스

end: 배열에서 현재 가리키는 오른쪽 인덱스

idx: 트리에서 현재 가리키는 인덱스

diff: 배열에서 값을 변경하는 위치의 인덱스

value: 배열에서 값을 변경하는 크기

 

현재 가리키는 구간 내에 변경되는 값의 위치가 포함되어있지 않다면 return 하고, 그렇지 않다면 트리의 현재 인덱스의 값을 변경값만큼 늘린다.

start와 end가 같다면 더 이상 분할할 것이 없으므로 그대로 return하고, 그렇지 않다면 mid를 기준으로 왼쪽 구간과 오른쪽 구간을 각각 탐색한다.

 

def update(start, end, idx, diff, value):
    if diff < start or end < diff:
        return

    tree[idx] += value
    if start == end:
        return

    mid = (start + end) // 2
    update(start, mid, idx*2, diff, value)
    update(mid+1, end, idx*2+1, diff, value)

 

 

특정 구간의 합 또는 그것을 이용하는 작업을 수행해야 하는데 배열의 크기가 너무 크다면, 누적 합 알고리즘 대신 이 세그먼트 트리 자료구조를 만들어 활용하면 훨씬 더 빠르게 원하는 결과를 찾을 수 있을 것이다.