세그먼트 트리

개요

선형적인 자료구조에서는 값에 접근하는 데에 \(O(1)\)이면 충분하지만, 대신 부분합을 구하는 데에는 \(O(N)\)이 필요하다.
그렇다면 이 자료구조를 이진 트리로 구성하면 어떨까? 값에 접근하는 데에 걸리는 시간이 \(O(\lg N)\)으로 늘어나지만 대신 부분합을 구하는 시간도 \(O(\lg N)\)으로 줄어들게 된다.
이러한 자료구조를 세그먼트 트리라고 한다.
그렇다면 세그먼트 트리는 어떻게 동작하는가?
세그먼트 트리의 리프 노드는 값 자체가 되고, 리프 노드가 아닌 노드는 자식 노드들의 합이 된다.
따라서 모든 노드의 조상인 루트 노드는 모든 노드의 합이 될 것이다.
특정 구간의 합은 특정 노드들의 합으로 정의되므로 부분합을 쉽게 구할 수 있다.

예시

예를 들어 1에서 4의 값이 있다고 하자.

image
리프 노드는 위와 같이 정의될 것이다.

image
리프 노드의 부모 노드들은 리프 노드의 합으로 정의된다.

image
마지막으로 그래프 전체는 위와 같아질 것이다.

image
만약 0번 노드(값: 1)부터 2번 노드(값: 3)의 합을 구한다면 색칠된 노드들을 더하면 될 것이다. 따라서 그 합은 6이다.

값을 업데이트하는 경우는 어떨까?

image
2번 노드(값: 3)의 값을 5로 변경하면 선형적 자료구조에서는 위와 같다. 아직 트리에 반영되지 않았음에 유의하자.

image
이제 색칠한 리프 노드가 값이 바뀐다.

image
하지만 여기서 끝내면 안 되는데, 그 부모 노드들 역시 값이 바뀌어야 하기 때문이다. 부모 노드가 5+4=9로 값이 바뀌었다.

image
마지막으로 루트 노드까지 값이 바뀌면 업데이트가 끝나게 된다.

구현

# 초기화
def init(node, l, r):
    if l == r:
        # 리프 노드인 경우
        segtree[node] = arr[l]
    else:
        # 그 외의 경우 왼쪽 자식과 오른쪽 자식의 합
        m = (l + r) // 2
        segtree[node] = init(node*2, l, m) + init(node*2+1, m + 1, r)
    return segtree[node]

# s부터 e까지의 합 구하기
def getSum(node, l, r, s, e):
    # l, r : node의 범위
    # s, e : 합을 구해야 하는 범위
    if l > e or r < s:
        # 두 범위가 겹치지 않을 경우
        return 0
    elif s <= l and r <= e:
        # 노드의 범위가 구하고자 하는 범위에 포함될 경우
        return segtree[node]
    m = (l + r) // 2
    return getSum(node*2, l, m, s, e) + getSum(node*2+1, m+1, r, s, e)

# i번 노드를 diff만큼 더하는 업데이트
def update(node, l, r, i, diff):
    if i < l or i > r:
        # i가 해당 범위에 없을 경우
        return
    # 그 외의 경우 값을 갱신
    segtree[node] += diff
    if l != r:
        # 리프 노드가 아니면 자식 노드도 갱신
        m = (l + r) // 2
        update(node*2, l, m, i, diff)
        update(node*2+1, m+1, r, i, diff)

arr = [i+1 for i in range(4)]
init(1, 0, n-1)

2024

맨 위로 이동 ↑

2023

세그먼트 트리

개요 선형적인 자료구조에서는 값에 접근하는 데에 \(O(1)\)이면 충분하지만, 대신 부분합을 구하는 데에는 \(O(N)\)이 필요하다. 그렇다면 이 자료구조를 이진 트리로 구성하면 어떨까? 값에 접근하는 데에 걸리는 시간이 \(O(\lg N)\)으로 늘어나지만 대신 부분합을 구하...

벨만-포드 알고리즘

개요 다익스트라 알고리즘과 함께 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 다익스트라 알고리즘보다 느리지만, 음수 가중치 간선이 있어도 작동하며, 음수 가중치 사...

다익스트라 알고리즘

개요 다익스트라 알고리즘은 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘 중 하나이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 단, 다익스트라 알고리즘은 음수 가중치 엣지를 허용하지 않는다. 이 경우에는 벨만-...

맨 위로 이동 ↑