[4주차] 세그먼트 트리 (Segment tree)

 

 

이번 주차에서는 자료구조를 응용하고 그래프 심화 알고리즘을 다룬다.

하나하나가 다뤄야 할 양이 방대해서... 공부 내용 정리는 한 가지씩 해보려고 한다.

오늘 다룰 내용은 세그먼트 트리 (Segment tree)다.

 

---

 

구간합(Prefix sum)을 구한다. (근데 값의 갱신이 있는...)

 

1차원 배열에서 구간합을 구한다고 생각해보자.

널리 알려진 구간합을 구하는 알고리즘이 있으니 어렵지 않게 구할 수 있다.

 

원본 배열을 A, 이를 바탕으로 만드는 구간합 배열을 S라고 할 때 1 ~ i 구간의 구간합 배열 S는 아래와 같이 구할 수 있다.

S[i] = S[i - 1] + A[i]

 

그리고 i ~ j 사이의 구간합을 구할 때에는

S[j] - S[i - 1]

 

위 처럼 구한다는 사실은 너무나도 널리 알려져있다.

 

이 방식으로 구하면 구간합 배열을 만드는 데에 O(n), 구간합을 구하는 데에 질의 당 O(1)의 시간 복잡도를 갖는다.

하지만 이 구간합에서 원본 배열의 값의 갱신이 발생한다면?

 

구간합에 대한 질의는 O(1)의 시간 복잡도로 답변할 수 있지만, Update query가 들어온다면 해당 위치 뒤의 모든 구간합에 변화가 필요하고, 이 경우 질의의 갯수를 M이라고 하면 최악의 시간 복잡도는 O(NM)에 달한다.

 

따라서 이런 값이 갱신되는 구간합에서는 '세그먼트 트리 (Segment tree)'를 사용해야 한다.

 

세그먼트 트리의 기본 아이디어

 

세그먼트 트리의 기본 아이디어는 구간 합과 동일하게 '큰 덩어리들의 합을 미리 계산하여 답을 구한다.' 이다.

세그먼트 트리는 완전 이진 트리(Complete binary tree)로 구현되며, 구간을 표현하기 위해 고르는 노드는 같은 깊이에 최대 2개만 존재해야한다. 

이를 자세히 설명하자면,

 

1. 세그먼트 트리는 완전 이진 트리(Complete binary tree)로 구현된다.

세그먼트 트리를 실제로 구현할 때에는 배열로, 인덱스로 관리된다. 따라서 마지막 레벨이 완전히 채워지지 않더라도 왼쪽부터 채워지는 구조를 갖기에 완전 이진 트리로 구현된다고 하는 것이다.

 

2. 구간을 표현하기 위해 고르는 노드는 같은 깊이에 최대 2개만 존재해야 한다.

여기서 말하는 "같은 깊이에 최대 2개" 라는 말은 "특정 구간 [S, E]를 표현할 때"의 규칙을 말한다. 즉, 세그먼트 트리의 전체적인 구조에 대한 이야기아 아닌 특정 구간의 합을 표현하기 위해 선택되는 노드들을 말하는 것이다.

 

배열이 [1, 2, 3, 4, 5, 6, 7, 8] 이라고 한다면, 세그먼트 트리의 루트 노드는 [1, 8]의 전체 구간을 나타낸다.

여기서 [3, 7]의 구간합을 구한다고 가정하면 우리는 세그먼트 트리에서 이 구간을 완전히 포함하는 최소 개수의 노드를 고른다.

아래는 전체 트리를 레벨 별로 나타낸 것으로 빨간색으로 하이라이팅 된 부분이 각 레벨에서 선택되는 노드다.

 

level 1 - [1, 8] (1번 노드 / root)

level 2 - [1, 4], [5, 8] (2, 3번 노드)

leve 3 - [1, 2], [3, 4], [5, 6], [7, 8] (4, 5, 6, 7번 노드)

level 4 - [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7], [8, 8] (8, 9, 10 , 11, 12, 13, 14, 15번 노드)

 

위 예시로 보았을 때, 특정 구간을 표현할 때 같은 깊이에 2개 이하의 노드만 선택된 것을 볼 수 있다.

세그먼트 트리의 구조상 한 구간 [S, E]를 표현하기 위해서는 각 레벨에서 겹치지 않는 최대 두 개의 구간으로만 나눌 수 있기 때문이다.

 

* 완전 이진 트리 (Complete binary tree)

• 마지막 레벨을 제외한 모든 노드가 채워져 있어야 한다. 마지막 레벨의 노드는 다 채워져 있을 수도 있고 아닐 수도 있다.
• 노드는 반드시 왼쪽 자식부터 오른쪽 자식 순으로 채워져야 한다.

 

세그먼트 트리를 활용한 구간합 구하기 

백준 2042번 '구간합 구하기' 문제(https://www.acmicpc.net/problem/2042)는 세그먼트 트리 자체를 구현해야 하는 문제다. 

이 문제를 풀어나가면서 세그먼트 트리를 python으로 구현하는 방법을 알아보자.

 

세그먼트 트리는 탑다운 (Top-down) 방식으로 재귀 구현했다.

이번 포스팅에서는 세그먼트 트리의 기본 골자만을 다루므로 바텁업 방식은 추후에 기회가 된다면 빼서 따로 작성하려고 한다.

 

세그먼트 트리 초기화

 

세그먼트 트리 초기화란 주어진 배열의 값을 통해 세그먼트 트리 배열을 구성하는 것을 말한다.

tree = [0] * (4 * n)

# 세그먼트 트리 초기화
def init(arr, tree, node, start, end):
  # 노드가 리프 노드인 경우 
  # e. g. [start, end] => [1, 1]인 경우
  if start == end:
    tree[node] = arr[start]
    return tree[node]
  
  # 현재 노드가 담당하는 구간을 정확히 반으로 나눔
  # start ~ mid : 왼쪽 자식, mid + 1 ~ end: 오른쪽 자식
  mid = (start + end) // 2

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

 

 

보통 세그먼트 트리의 크기를 정할 때에는 간단하게 안전한 충분한 크기를 잡는 근사값 4 * n으로 초기화를 한다. 계산하기 귀찮기도 하고, 4* n은 최대 트리 크기를 항상 보장하기 때문이다.

 

* 정확한 세그먼트 트리의 노드 갯수 구하기

세그먼트 트리는 완전 이진 트리 형태로 만들어지기 때문에 리프 노드 수가 n일 때, 전체 노드 개수는 2 x (트리의 높이) - 1 이 된다.

 

트리의 높이는 루트에서 리프까지의 레벨 수로, 리프 노드가 n개 라면 다음과 같이 계산한다.

$$h = \lceil \log_2(n) \rceil + 1$$

 

import math

n = 5
height = math.ceil(math.log2(n)) + 1
tree_size = 2 ** height - 1

 

 

세그먼트 트리 구간합

 

구간합을 구한다.

# 구간합 구하기
def prefixSum(tree, node, start, end, left, right):
  # [left, right]의 구간합을 구한다.

  # [start, end] 앞 뒤에 [left, right]가 있는 경우 겹치지 않으므로 탐색할 필요가 없다.
  # [start, end] : 현재 노드가 포함하는 구간
  if left > end or right < start : 
    return 0
  
  #[start, end]가 [left, right]에 포함되는 경우
  if left <= start and end <= right:
    return tree[node]
  
  # 일부만 겹치는 경우 왼쪽 자식과 오른쪽 자식을 루트로 하는 하위 트리에서 재탐색
  mid = (start + end) // 2

  prefix = prefixSum(tree, node * 2, start, mid, left, right) + prefixSum(tree, node * 2 + 1, mid + 1, end, left, right)

  return prefix

 

세그먼트 트리 구간합 갱신

# 구간합 갱신
def update(tree, node, start, end, index, diff):
  # 현재 노드의 구간이 인덱스를 포함하지 않는다면 해당 하위 트리로 내려갈 필요가 없다.
  if index < start or index > end:
    return
  
  # 해당하면 차 만큼만 더해주면 된다.
  tree[node] = tree[node] + diff

  if start != end:
    mid = (start+ end) // 2
    # 왼쪽 자식, 오른쪽 자식 갱신
    update(tree, node * 2, start, mid, index, diff)
    update(tree, node * 2 + 1, mid + 1, end, index, diff)

 

전체 코드

import sys
from collections import deque
from collections import defaultdict

input = sys.stdin.readline

n, m, k = map(int, input().split())

arr = [0] * (n + 1)
tree = [0] * (4 * n)

# 세그먼트 트리 초기화
def init(arr, tree, node, start, end):
  # 노드가 리프 노드인 경우 
  # e. g. [start, end] => [1, 1]인 경우
  if start == end:
    tree[node] = arr[start]
    return tree[node]
  
  # 현재 노드가 담당하는 구간을 정확히 반으로 나눔
  # start ~ mid : 왼쪽 자식, mid + 1 ~ end: 오른쪽 자식
  mid = (start + end) // 2

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

# 구간합 구하기
def prefixSum(tree, node, start, end, left, right):
  # [left, right]의 구간합을 구한다.

  # [start, end] 앞 뒤에 [left, right]가 있는 경우 겹치지 않으므로 탐색할 필요가 없다.
  # [start, end] : 현재 노드가 포함하는 구간
  if left > end or right < start : 
    return 0
  
  #[start, end]가 [left, right]에 포함되는 경우
  if left <= start and end <= right:
    return tree[node]
  
  # 일부만 겹치는 경우 왼쪽 자식과 오른쪽 자식을 루트로 하는 하위 트리에서 재탐색
  mid = (start + end) // 2

  prefix = prefixSum(tree, node * 2, start, mid, left, right) + prefixSum(tree, node * 2 + 1, mid + 1, end, left, right)

  return prefix

# 구간합 갱신
def update(tree, node, start, end, index, diff):
  # 현재 노드의 구간이 인덱스를 포함하지 않는다면 해당 하위 트리로 내려갈 필요가 없다.
  if index < start or index > end:
    return
  
  # 해당하면 차 만큼만 더해주면 된다.
  tree[node] = tree[node] + diff

  if start != end:
    mid = (start+ end) // 2
    # 왼쪽 자식, 오른쪽 자식 갱신
    update(tree, node * 2, start, mid, index, diff)
    update(tree, node * 2 + 1, mid + 1, end, index, diff)

  


for i in range(1, n + 1):
  arr[i] = int(input())

# 세그먼트 트리 초기화
init(arr, tree, 1, 1, n)

for _ in range(m + k):
  a, b, c = map(int, input().split())

  if a == 1:
    # a가 1이면 갱신 (update)
    diff = c - arr[b]
    arr[b] = c
    update(tree, 1, 1, n, b, diff)
  elif a == 2:
    print(prefixSum(tree, 1, 1, n, b, c))

 

 

세그먼트 트리의 다른 활용

 

이 포스팅에서는 세그먼트 트리로 구간합을 구하는 데에 초점을 맞췄지만, 세그먼트 트리는 구간을 효율적으로 다뤄야 하는 모든 문제에서 사용할 수 있다. 저번에 알아봤던 모노토닉 큐라는 자료구조와 마찬가지로 슬라이딩 윈도우 최적화에도 사용될 수 있는 것이다.

 

세그먼트 트리의 핵심 아이디어는 결국 "배열의 일부분에 대한 정보를 빠르게 갱신하고 빠르게 질의하기 위한 자료구조"이므로, 구간합 이외에도 구간의 최솟값, 구간의 최댓값 등의 문제에도 사용될 수 있다.

 

다음에는 모노토닉 큐를 활용했던 문제들도 세그먼트 트리로 재풀이를 해보는 시간을 가지면 좋을 것 같다.

 

참고자료 출처

- segment tree (IOI KOREA) : https://www.youtube.com/watch?v=pINRhSfISKE

- segment tree (geeks for geeks) : https://www.geeksforgeeks.org/dsa/segment-tree-data-structure/

- 세그먼트 트리 (다빈치코딩 알고리즘 / 위키독스) : https://wikidocs.net/209446

- 세그먼트 트리 (나무위키) : https://namu.wiki/w/%EC%84%B8%EA%B7%B8%EB%A8%BC%ED%8A%B8%20%ED%8A%B8%EB%A6%AC