세그먼트 트리

지연 전파(lazy propagation)에 대한 예제는 이 글을 참고한다.

세그먼트 트리를 소개한다. 이 자료구조는 매우 유연하고 구현이 간단하여 내가 가장 좋아하는 자료구조 중 하나다.

어떤 타입의 배열 a와 결합 법칙이 성립하는 함수 f가 있다고 가정한다. 예를 들어, 이 함수는 합, 곱, 최솟값, 최댓값, 최대공약수 등이 될 수 있다.

여러분의 과제는 다음과 같다:

• l과 r로 주어진 구간에 대한 쿼리에 답한다. 즉, f(a[l], a[l+1], ..., a[r-1], a[r])을 수행한다.
• 특정 인덱스의 값을 교체하는 기능을 지원한다. 예: a[index] = newItem

예를 들어, 숫자 배열이 다음과 같다면:

var a = [ 20, 3, -1, 101, 14, 29, 5, 61, 99 ]

이 배열에서 3부터 7까지의 구간에 대해 "합" 함수를 쿼리한다고 하자. 이 경우 다음과 같은 계산을 수행한다:

101 + 14 + 29 + 5 + 61 = 210

배열에서 101은 인덱스 3에 위치하고, 61은 인덱스 7에 위치한다. 따라서 101부터 61까지의 모든 숫자를 합 함수에 전달하여 더한다. 만약 "최솟값" 함수를 사용했다면 결과는 5가 된다. 이는 3부터 7까지의 구간에서 가장 작은 숫자이기 때문이다.

배열의 타입이 Int이고 f가 단순히 두 정수의 합이라면, 다음과 같이 간단한 방법으로 구현할 수 있다:

func query(array: [Int], l: Int, r: Int) -> Int {
  var sum = 0
  for i in l...r {
    sum += array[i]
  }
  return sum
}

이 알고리즘의 시간 복잡도는 최악의 경우 **O(n)**이다. 즉, l = 0, r = n-1인 경우(여기서 n은 배열의 요소 수)에 해당한다. 만약 m개의 쿼리를 처리해야 한다면 전체 복잡도는 **O(m*n)**이 된다.

배열의 크기가 100,000(n = 10^5)이고 100개의 쿼리(m = 100)를 처리해야 한다면, 이 알고리즘은 10^7번의 작업을 수행한다. 이는 매우 비효율적이다. 이를 개선하는 방법을 살펴보자.

세그먼트 트리를 사용하면 쿼리에 답하고 항목을 교체하는 작업을 O(log n) 시간에 처리할 수 있다. 마법 같지 않은가? ✨

세그먼트 트리의 핵심 아이디어는 간단하다. 배열의 일부 구간을 미리 계산해 두고, 이를 반복적으로 사용하는 것이다.

세그먼트 트리의 구조

세그먼트 트리는 기본적으로 이진 트리의 한 형태로, 각 노드가 SegmentTree 클래스의 인스턴스로 구성된다:

public class SegmentTree<T> {
  private var value: T
  private var function: (T, T) -> T
  private var leftBound: Int
  private var rightBound: Int
  private var leftChild: SegmentTree<T>?
  private var rightChild: SegmentTree<T>?
}

각 노드는 다음과 같은 데이터를 가지고 있다:

• leftBound와 rightBound는 특정 구간을 나타낸다.
• leftChild와 rightChild는 자식 노드를 가리키는 포인터다.
• value는 함수 f(a[leftBound], a[leftBound+1], ..., a[rightBound-1], a[rightBound])를 적용한 결과값이다.

예를 들어, 배열이 [1, 2, 3, 4]이고 함수 f = a + b일 때, 세그먼트 트리는 다음과 같은 구조를 가진다:

각 노드의 leftBound와 rightBound는 빨간색으로 표시되어 있다.

세그먼트 트리 구축

세그먼트 트리의 노드를 생성하는 방법은 다음과 같다:

public init(array: [T], leftBound: Int, rightBound: Int, function: @escaping (T, T) -> T) {
    self.leftBound = leftBound
    self.rightBound = rightBound
    self.function = function

    if leftBound == rightBound {                    // 1
      value = array[leftBound]
    } else {
      let middle = (leftBound + rightBound) / 2     // 2

      // 3
      leftChild = SegmentTree<T>(array: array, leftBound: leftBound, rightBound: middle, function: function)
      rightChild = SegmentTree<T>(array: array, leftBound: middle+1, rightBound: rightBound, function: function)

      value = function(leftChild!.value, rightChild!.value)  // 4
    }
  }

이 코드는 재귀적 메서드이다. [1, 2, 3, 4]와 같은 배열을 입력하면, 루트 노드부터 모든 자식 노드까지 전체 트리를 구축한다.

1. leftBound와 rightBound가 같으면 재귀가 종료된다. 이런 SegmentTree 인스턴스는 리프 노드를 나타낸다. 입력 배열 [1, 2, 3, 4]의 경우, 1, 2, 3, 4 네 개의 리프 노드가 생성된다. 배열의 숫자를 value 속성에 그대로 할당한다.

2. rightBound가 leftBound보다 크면 두 개의 자식 노드를 생성한다. 현재 세그먼트를 두 개의 동일한 세그먼트로 나눈다(길이가 짝수일 경우; 홀수일 경우 하나의 세그먼트가 약간 더 커진다).

3. 두 세그먼트에 대해 자식 노드를 재귀적으로 구축한다. 왼쪽 자식 노드는 [leftBound, middle] 구간을, 오른쪽 자식 노드는 [middle+1, rightBound] 구간을 담당한다.

4. 자식 노드를 생성한 후, f(leftBound, rightBound) = f(f(leftBound, middle), f(middle+1, rightBound)) 공식을 이용해 자신의 값을 계산한다. 이는 수학적 원리이다!

트리를 구축하는 작업은 O(n) 시간 복잡도를 가진다.

쿼리에 대한 답 얻기

이 모든 과정은 트리를 효율적으로 쿼리하기 위해 진행한다.

다음은 관련 코드다:

public func query(withLeftBound: leftBound: Int, rightBound: Int) -> T {
    // 1
    if self.leftBound == leftBound && self.rightBound == rightBound {
      return self.value
    }

    guard let leftChild = leftChild else { fatalError("leftChild should not be nil") }
    guard let rightChild = rightChild else { fatalError("rightChild should not be nil") }

    // 2
    if leftChild.rightBound < leftBound {
      return rightChild.query(withLeftBound: leftBound, rightBound: rightBound)

    // 3
    } else if rightChild.leftBound > rightBound {
      return leftChild.query(withLeftBound: leftBound, rightBound: rightBound)

    // 4
    } else {
      let leftResult = leftChild.query(withLeftBound: leftBound, rightBound: leftChild.rightBound)
      let rightResult = rightChild.query(withLeftBound: rightChild.leftBound, rightBound: rightBound)
      return function(leftResult, rightResult)
    }
  }

이 메서드는 재귀적으로 동작하며, 네 가지 가능성을 확인한다.

1. 먼저, 쿼리 구간이 현재 노드가 담당하는 구간과 동일한지 확인한다. 동일하다면 현재 노드의 값을 반환한다.

2. 쿼리 구간이 오른쪽 자식에 완전히 포함되는지 확인한다. 포함된다면 오른쪽 자식에 대해 재귀적으로 쿼리를 수행한다.

3. 쿼리 구간이 왼쪽 자식에 완전히 포함되는지 확인한다. 포함된다면 왼쪽 자식에 대해 재귀적으로 쿼리를 수행한다.

4. 위의 경우가 모두 아니라면, 쿼리가 두 자식에 부분적으로 걸쳐 있음을 의미한다. 따라서 두 자식에 대한 쿼리 결과를 결합한다.

다음은 플레이그라운드에서 테스트하는 방법이다:

let array = [1, 2, 3, 4]

let sumSegmentTree = SegmentTree(array: array, function: +)

sumSegmentTree.query(withLeftBound: 0, rightBound: 3)  // 1 + 2 + 3 + 4 = 10
sumSegmentTree.query(withLeftBound: 1, rightBound: 2)  // 2 + 3 = 5
sumSegmentTree.query(withLeftBound: 0, rightBound: 0)  // just 1
sumSegmentTree.query(withLeftBound: 3, rightBound: 3)  // just 4

트리를 쿼리하는 데 걸리는 시간은 **O(log n)**이다.

아이템 교체

세그먼트 트리에서 노드의 값은 하위 노드에 의존한다. 따라서 리프 노드의 값을 변경하려면 모든 상위 노드도 함께 업데이트해야 한다.

다음은 관련 코드다:

  public func replaceItem(at index: Int, withItem item: T) {
    if leftBound == rightBound {
      value = item
    } else if let leftChild = leftChild, rightChild = rightChild {
      if leftChild.rightBound >= index {
        leftChild.replaceItem(at: index, withItem: item)
      } else {
        rightChild.replaceItem(at: index, withItem: item)
      }
      value = function(leftChild.value, rightChild.value)
    }
  }

이 코드는 재귀를 사용한다. 노드가 리프 노드라면 값을 변경한다. 리프 노드가 아니라면 replaceItem(at: )을 재귀적으로 호출해 자식 노드를 업데이트한다. 이후 노드의 값을 다시 계산해 최신 상태로 유지한다.

아이템을 교체하는 데 O(log n) 시간이 걸린다.

세그먼트 트리를 사용하는 더 많은 예제는 플레이그라운드를 참고한다.

관련 자료

레이지 프로퍼게이션(Lazy Propagation) 구현 및 설명.

PEGWiki의 세그먼트 트리(Segment tree)

Swift Algorithm Club을 위해 Artur Antonov이 작성함