세그먼트 트리

세그먼트 트리가 필요한 상황

세그먼트 트리는 [1,2,3,4,5,6,7,8,9]와 같은 수열이 있을 때, 주어진 수열의 부분 수열들에서 특정 값(ex. 최댓값, 최솟값, 총 합)을 빠르게 알고 싶을 때 사용하는 자료구조다.

예를 들어, [1,2,3,4,5,6,7,8,9] 수열의 두번째부터 다섯번째까지의 부분수열 총합을 구하면 2 + 3 + 4 + 5 = 14다.

이 값을 구하기 위해서는 두번째부터 다섯번째까지 모두 값을 하나하나 더해야 한다. 부분집합의 길이가 m이라고 한다면 시간복잡도는 O(m).

이러한 연산을 자주 하는 상황이 발생할 때 매번 부분집합의 길이만큼 연산을 하는것은 비효율적이다.

일반적으로 자주 발생하는 연산을 줄이는 방법은 자주 발생하는 연산의 결과를 저장해 놓는것이다.

하지만 n길이의 수열에서 부분 수열의 개수는 n + (n-1) + (n-2) + ... + 1 = n(n+1) / 2로 n이 큰 값일 때는 메모리 공간이 부족할 수 있다.

세그먼트 트리는 이런 시간복잡도와 공간복잡도의 상충관계를 적절히 조절할 수 있는 자료구조다.

세그먼트 트리의 구조

우선 부분 수열의 합을 알고 싶은 상황이라 가정한다.

수열을 두 개씩 묶고 해당 부분 수열에서 원하는 값을 추출해서 부모 노드로 만든다. 이를 전체 수열에 대한 값이 나올 때 까지 반복하면 이진 트리가 생성된다.

두 개씩 묶는것이 유일한 방법은 아니지만 두개 씩 묶는것이 구현도 쉽고 성능도 충분하다.

이제 생성된 세그먼트 트리에서 원하는 부분 수열의 값을 구해보자. 두번째부터 여섯번째까지의 부분수열 합을 구하려면

2 + 7(3+4) + 11(5+6) = 20, 원래 대로라면 2, 3, 4, 5, 6 다섯개의 숫자 합에서 세개의 노드 합으로 연산을 줄일 수 있다.

예시의 수열 크기가 작아 크게 효과적으로 보이지 않을 수 있지만 N이 커질 수록 연산의 수가 꽤나 감소하게 된다.

구현

트리 선언

h = math.ceil(math.log2(n)) + 1
size = 1 << h
segment = [0 for _ in range(size)]

이진 트리의 높이는 log2(n)이고 이진 트리의 최대 노드 개수는 2^n- 1다.

트리 초기화

def init(idx, s, e):
    if s == e:
        segment[idx] = arr[idx]
        return segment[idx]
    mid = (s + e) // 2
    l = init(idx * 2, s, mid)
    r = init(idx * 2 + 1, mid + 1, e)
    segment[idx] = l + r
    return segment[idx]

이진 트리의 왼쪽과 오른쪽을 재귀적으로 탐색해가면서 세그먼트 트리의 각 노드 값을 초기화 해준다.

idx는 각 노드번호라고 생각하면 된다. 이진트리에서 왼쪽 자식으로 갈 때는 자신의 노드번호의 두배, 오른쪽 자식으로 갈 때는 자신의 노드번호의 두배 + 1로 접근할 수 있다.

부분 수열의 값 탐색

def find(s, e, left, right, idx):
    if e < left or right < s: return 0
    mid = (s + e) // 2
    if left <= s and e <= right:
        return segment[idx]
    else:
        l = find(s, mid, left, right, idx * 2)
        r = find(mid+1, e, left, right, idx * 2 + 1)
        return l + r

#find(0, n, left, right, 1)와 같이 사용하고 left 부터 right 까지의 부분수열 값을 반환한다.