[417] 세그먼트 트리 파이썬 코드

1. 세그먼트 트리 개념

이번 글에서는 세그먼트 트리를 파이썬 코드로 어떻게 구현할 수 있는지 알아봅시다.

세그먼트 트리의 개념이 궁금하신 분은 지난 글을 참고해주세요.

https://kimwooil.tistory.com/416

[416] 세그먼트 트리의 개념

 

지난 글에서 예시로 다루었던 1~10까지의 누적합을 담는 세그먼트 트리를 파이썬 코드로 구현해보겠습니다.

트리를 다시 보면, 노드들이 다음과 같은 범위의 정보를 가집니다.

해당 범위의 누적합(구간합)은 다음과 같이 그려집니다.

2. 변수 선언

세그먼트 트리를 만들기 위해서는 원소가 되는 배열(arr)과 완전 이진 트리를 만들 배열(tree)이 필요합니다.

tree의 길이는 원소 개수에 4를 곱합니다.(자세한 설명은 세그먼트 트리의 개념을 다룬 글을 참고해주세요.)

arr = [0]+[*range(1, 10+1)]
tree = [0]*10*4

3. 초기화 함수

초기 세그먼트 트리를 만들 함수가 필요합니다.

완전 이진 트리이므로 계속해서 범위를 반으로 줄여가며 해당 노드에 재귀적으로 반환되는 값을 담습니다.

인덱스 값은 왼쪽 자식 노드의 경우 2배, 오른쪽 자식 노드의 경우 2배에 1을 더합니다.

def init(start, end, i):
    if start == end:
        tree[i] = arr[start]
    else:
        mid = (start+end)//2 # 범위가 반으로 줄어듭니다.
        # 왼쪽 자식 노드는 인덱스 값이 2배, 오른쪽 자식 노드는 2배에 1을 더한 값입니다.
        tree[i] = init(start, mid, i*2) + init(mid+1, end, i*2+1)
    return tree[i]

실행하면 tree 배열에 세그먼트 트리가 만들어집니다.

init(1, 10, 1)

print(tree)
# [0, 55, 15, 40, 6, 9, 21, 19, 3, 3, 4, 5, 13, 8, 9, 10, 1, 2, 0, 0, 0, 0, 0, 0, 6, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

4. 특정 구간의 값을 구하는 함수

역시 재귀적으로 범위를 나눠가며 탐색합니다.

먼저, 해당 노드가 주어진 범위에 전혀 포함되지 않으면 0을 반환합니다.

다음으로, 해당 노드가 주어진 범위에 완전히 들어간다면 노드의 값을 반환합니다.

마지막으로 해당 노드가 주어진 범위의 일부를 포함한다면 범위를 반으로 나누어 자식 노드에게 넘깁니다.

def query(start, end, i, left, right): # start, end, i는 세그먼트 트리의 노드 정보입니다.
    if left > end or right < start: # 원하는 범위(left ~ right)가 노드의 범위를 벗어난 경우입니다.
        return 0
    if left <= start and end <= right: # 노드의 범위가 원하는 범위에 포함되는 경우입니다.
        return tree[i]
    # 그 외에는 일부만 포함되는 경우이므로, 범위를 반으로 나누어 자식 노드에게 넘깁니다.
    mid = (start+end)//2
    return query(start, mid, i*2, left, right) + query(mid+1, end, i*2+1, left, right)

함수를 실행하면 위 그림과 같은 과정을 거쳐 3 ~ 7 범위의 구간합인 25가 출력됩니다.

result = query(1, 10, 1, 3, 7)

print(result)
# 25

5. 원소의 값을 수정하는 함수

두 가지 방법이 존재합니다.

먼저, 부모에서부터 자식으로 내려가면서 바로바로 노드의 값을 수정하는 경우입니다.

이 경우는 원소의 변화량을 인자로 받습니다.

유효한 노드의 경우에만 변화량을 더해가며 노드를 수정합니다.

def update(start, end, i, n, diff):
    if n < start or n > end:
        return
    tree[i] += diff
    if start == end:
        return
    mid = (start+end)//2
    update(start, mid, i*2, n, diff)
    update(mid+1, end, i*2+1, n, diff)

5번 원소의 값을 10으로 수정하고 싶다면, 변화량인 5를 인자로 넘깁니다.

실행 후 tree 배열을 출력하면 위 그림의 세그먼트 트리와 같이 수정된 결과를 확인할 수 있습니다.

update(1, 10, 1, 5, 5) # 5번 원소의 값을 5 증가시킵니다.

print(tree)
# [0, 60, 20, 40, 6, 14, 21, 19, 3, 3, 4, 10, 13, 8, 9, 10, 1, 2, 0, 0, 0, 0, 0, 0, 6, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

 

두번째로 재귀적으로 반환값을 가지고 올라오면서 노드의 값을 수정하는 함수입니다.

유효한 범위가 아닌 경우, 단순히 노드의 값을 반환합니다.

수정할 원소에 도달한 경우 원소의 값을 수정하여 반환합니다.

반환된 값을 가지고 부모 노드의 값을 수정합니다.

def update(start, end, i, n, k):
    if n < start or n > end:
        pass
    elif start == end:
        if start == n: # 원하는 노드에 도착한 경우입니다.
            tree[i] = k
    else:
        mid = (start+end)//2
        tree[i] = update(start, mid, i*2, n, k) + update(mid+1, end, i*2+1, n, k)
    return tree[i]

함수를 실행할 때, 수정할 값을 인자로 넘기면 위 그림과 같은 과정을 거쳐 세그먼트 트리가 수정됩니다.

update(1, 10, 1, 5, 10) # 5번 원소의 값을 10으로 수정합니다.

print(tree)
# [0, 60, 20, 40, 6, 14, 21, 19, 3, 3, 4, 10, 13, 8, 9, 10, 1, 2, 0, 0, 0, 0, 0, 0, 6, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]