[파이썬으로 구현하는] 세그먼트 트리(Segment Tree) - Top-Down

---

"파이썬으로 구현하는"

시리즈 2

세그먼트 트리

(Segment Tree)

Top-Down

---

오늘은 파이썬으로 구현하는 세그먼트 트리에 대해서 알아볼 것이다.

 

먼저 세그먼트 트리가 무엇인지 알아보자.

 

가끔씩 우리는 어떤 수열에서 연속한 부분에 대한 쿼리를 처리해주어야 할 때가 있다.

예를 들어, [3, 1, 4, 1, 5, 9, 2] 라는 수열이 있다고 하자. 인덱스는 0부터 시작한다.

그리고, 인덱스 a부터 b까지 구간합을 구하는 쿼리가 주어진다.

여기까지는 딱히 새로운 자료구조 없이 누적합을 이용하면 쿼리가 아무리 많아도 쉽게 구할 수 있을 것이다.

 

그런데, 이번에는 저 리스트의 특정 인덱스의 값을 다른 값으로 변화시키는 쿼리가 들어왔다.

이런 경우에는 일일이 누적 합 배열을 다시 구하고 있어야 한다!

 

이런 경우(구간 쿼리, 값 업데이트 쿼리)에 쓸 수 있는 자료 구조가 바로 세그먼트 트리이다.

 

아이디어는 다음과 같다.

 

트리를 만드는데, 리프 노드들이 수열의 값을 가지게 한다.

그리고, 두 노드에 대해서 부모 노드는 자식 노드들의 합을 저장하고 있는다.

 

말로만 하면 이해가 안되니까 아래 그림을 살펴보자.

 

3, 1, 4, 1, 5, 9, 2의 값들을 리프 노드에 넣어주었다.

우리는 여기서 구간의 합을 빠르게 구해줄 수 있다.

예를 들어, 인덱스 1부터 4까지의 합을 구해준다고 하자.

구하는 합은 1+4+1+5 = 11인데, 이를 어떻게 구할까?

빨간색으로 표시된 노드들만 더해주면 될 것이다.

이는 인덱스 2 ~ 3 구간의 경우에는 미리 구해놓은 값(5)가 있기 때문이다.

이 빨간색 노드들을 찾는 것이 세그먼트 트리 구현의 핵심이라고 할 수 있다.

 

만약 값이 업데이트 된다면 어떻게 할까?

예를들어, 인덱스 2번의 값이 4에서 3으로 바뀐다고 하자.

우리는 인덱스 2번이 원래 포함되었었던 구간합 노드들의 값을 차례차례로 변경해주면 된다.

그림으로 나타내면 아래와 같다.

 

기존 값(4)와 바뀐 값(3)의 차이가 -1이라서, 해당 노드들에 -1를 더하는 방식으로 구현할 수 있다.

 

Top-Down 방식의 세그먼트 트리는 재귀로 구현하는데, 우리가 구현할 함수는 다음과 같다.

1. build() - 맨 처음 트리를 만드는 함수

2. search() - 구간합 쿼리

3. update() - 값 변경

 

일단, 맨 처음에 트리 리스트를 만들어주어야 한다.

트리의 size는 다음과 같이 찾을 수 있다.

arr = [3, 1, 4, 1, 5, 9, 2]
N = len(arr)

size = 1
while size < N:
    size <<= 1
tree = [0] * (2*size)
print(len(tree)) # 16

 

우리는 원소를 7개 가지고 있다.

그러면 리프 노드가 최소 7개 필요하고, 가장 가까우면서도 큰 2의 거듭제곱인 8개의 리프노드를 만들어준다.

그리고 리프노드들의 부모 노드들이  8 - 1 개 필요하다. 그래서 tree를 코드와 같이 만들어준다.

 

다음은 build() 함수이다.

def build(s, e, i=1): #i는 트리에서 인덱스
    if s == e:
        tree[i] = arr[s]
        return tree[i]
    mid = (s + e) // 2
    tree[i] = build(s, mid, i*2) + build(mid+1, e, i*2+1)
    return tree[i]

 

i는 tree 리스트 상의 인덱스이다.

s, e(start, end)는 그 노드가 포함하고 있는 범위이다. 당연히 처음에는 0, N-1이 들어간다.

가면서 s, e는 점점 좁혀지는데, 그러다가

s와 e가 같을 때, 즉 해당되는 구간이 노드 하나밖에 없을 때는 그 노드는 리프 노드라는 뜻이므로, arr에서 직접 값을 지정해준다.

그렇지 않은 경우에는 재귀를 이용해준다.

build(0, N-1) #s, e는 arr 리스트의 인덱스 기준! (inclusive)
print(tree) # [0, 25, 9, 16, 4, 5, 14, 2, 3, 1, 4, 1, 5, 9, 0, 0]

 

다음으로 search() 함수를 구현하자.

이 함수에서는 s, e, 그리고 구간 l, r(left, right), 트리상의 인덱스 i, 총 5개의 매개변수를 받는다.

def search(s, e, l, r, i = 1):
    #s, e는 i 노드가 포함하고 있는 범위 inclusive
    #l, r는 질의의 구간 inclusive
    #i는 여전히 트리에서의 인덱스
    #여기서 좁혀지는 범위는 s와 e이다. s와 e가 l, r에 맞추어지는 것.
    #즉, l, r는 바뀌지 않는다!

    if e < l or s > r: #찾는 범위 아닐 때
        return 0

    if l <= s and e <= r: #해당 범위에 딱 포함될 때
        return tree[i]

    mid = (s + e) // 2
    return search(s, mid, l, r, i*2) + search(mid+1, e, l, r, i*2+1)

s, e가 범위에 맞지 않을때, 즉 i 노드가 구간 질의에 하나도 포함되지 않는 범위를 가지고 있는 경우는 고려할 필요가 없으므로 0을 리턴한다.

반대로 i 노드가 포함하고 있는 범위(s, e)가 구간 질의(l, r)에 딱 맞게 완전히 포함되는 경우에는 그냥 그 노드를 반환하면 될 것이다.

그렇지 않은 경우는 반으로 나누어서 재귀를 써준다.

print(search(0, 6, 1, 4)) # 11

 

update() 함수를 살펴보자.

이번에는 s, e, 리스트에서 바꿀 idx, 기존 원소와의 차이인 diff, 그리고 i를 매개변수로 받는다.

def update(s, e, idx, diff, i=1):
    #idx는 변하고자 하는 리스트의 인덱스(트리인덱스 아님)
    #트리인덱스는 i에서 다룬다.

    if s > idx or idx > e:
        return
        
    tree[i] += diff
    
    if s != e:
        mid = (s + e) // 2
        update(s, mid, idx, diff, i*2)
        update(mid+1, e, idx, diff, i*2 + 1)

 

만약 해당 범위가 아니라면 (s > idx, idx > e)그냥 리턴한다.

나머지 경우에는 diff만큼 더해주고, 재귀한다.

update(0, 6, 2, -1)
print(tree) #[0, 24, 8, 16, 4, 4, 14, 2, 3, 1, 3, 1, 5, 9, 0, 0]

 

구현을 완성했다!

전체 코드를 모아보자.

 

그럼 이제 최종적인 구현 코드들을 정리해보자.

 

arr = [3, 1, 4, 1, 5, 9, 2]
N = len(arr)

size = 1
while size < N:
    size <<= 1
tree = [0] * (2*size)

def build(s, e, i=1): #i는 트리에서 인덱스
    if s == e:
        tree[i] = arr[s]
        return tree[i]
    mid = (s + e) // 2
    tree[i] = build(s, mid, i*2) + build(mid+1, e, i*2+1)
    return tree[i]

def search(s, e, l, r, i = 1):
    if e < l or s > r: #찾는 범위 아닐 때
        return 0
    
    if l <= s and e <= r: #해당 범위에 딱 포함될 때
        return tree[i]

    mid = (s + e) // 2
    return search(s, mid, l, r, i*2) + search(mid+1, e, l, r, i*2+1)

def update(s, e, idx, diff, i=1):
    if s > idx or idx > e:
        return

    tree[i] += diff

    if s != e:
        mid = (s + e) // 2
        update(s, mid, idx, diff, i * 2)
        update(mid + 1, e, idx, diff, i * 2 + 1)

build(0, N-1)
print(tree)

print(search(0, 6, 1, 4))

update(0, 6, 2, -1)
print(tree)

 

 

대표적인 연습문제로는

 

구간 합 구하기

https://www.acmicpc.net/problem/2042

 

최솟값과 최댓값

https://www.acmicpc.net/problem/2357

 

사탕상자

https://www.acmicpc.net/problem/2243

 

등이 있다.