[자료구조/Python] 세그먼트 트리(Segment Tree)

세그먼트 트리(Segment Tree)

구간 합, 최소값, 최대값 등의 쿼리를 빠르게 처리하기 위해 사용되는 자료구조 
이진 트리(binary tree)를 기반으로 하며, 각 노드는 해당 구간의 합, 최소값, 최대값 등의 값을 저장한다.

https://yoongrammer.tistory.com/103

1. 구성

1. 기본 배열을 이용하여 세그먼트 트리를 초기화, 배열의 각 요소가 세그먼트 트리의 각 리프 노드에 대응된다.
2. 이진 트리의 성질을 이용하여, 각 노드가 나타내는 구간을 두 개의 서브트리로 나눈다.
3. 각 노드에 대해, 해당 구간의 합, 최소값, 최대값 등의 값을 계산하여 저장한다.
4. 각 쿼리에 대해, 세그먼트 트리를 탐색하여 필요한 값을 계산한다.

리프 노드 : 배열의 그 수 자체
다른 노드 : 왼쪽 자식과 오른쪽 자식의 합을 저장

 

2. 시간 복잡도

트리의 높이가 log(N)이므로, 쿼리 처리 시간은 O(log N)

구간 합, 최소값, 최대값 외에도 다른 연산(예: 구간 최소값의 인덱스 찾기)을 빠르게 처리할 수 있는 장점

 

3. 장점

• 구간 합, 구간 최소값, 구간 최대값 등을 빠르게 계산할 수 있다.
• 쿼리에 대해 O(log n)의 시간 복잡도로 처리할 수 있다.
• 노드 수가 데이터 크기에 비례하므로 메모리 사용량이 적다.

 

4. 단점

• 초기화 시간이 O(n), 세그먼트 트리를 초기화하는 데는 O(n log n)의 시간이 소요됨.
• 각 노드에 저장하는 값이 구간의 요약 정보이기 때문에, 구간의 값이 변경될 때마다 해당 구간을 포함하는 모든 노드의 값을 업데이트해주어야 한다. 이로 인해 구간 업데이트의 시간 복잡도가 O(log n)

 

5. 활용 예시

• 구간 합 구하기
  • 주어진 배열의 특정 구간 내의 합을 구하는 문제
• 구간 최솟값, 최댓값 구하기
  • 주어진 배열의 특정 구간 내의 최솟값이나 최댓값을 구하는 문제
• 구간 업데이트
  • 주어진 배열의 특정 구간에 대해 값을 업데이트하는 문제
• 수열의 구간에 대한 통계값 구하기
  • 주어진 수열의 특정 구간에 대해 중앙값, 평균값, 표준편차 등의 통계값을 구하는 문제
• 최대 연속 부분합 구하기
  • 주어진 배열에서 최대 연속 부분합을 구하는 문제

 

6. 파이썬 구현 코드

class SegmentTree:
    def __init__(self, arr):
        self.arr = arr
        self.tree = [0] * (4 * len(arr))
        self._build(1, 0, len(arr) - 1)

    def _build(self, node, start, end):
        if start == end:
            self.tree[node] = self.arr[start]
        else:
            mid = (start + end) // 2
            self._build(2 * node, start, mid)
            self._build(2 * node + 1, mid + 1, end)
            self.tree[node] = self.tree[2 * node] + self.tree[2 * node + 1]

    def update(self, index, value):
        diff = value - self.arr[index]
        self.arr[index] = value
        self._update(1, 0, len(self.arr) - 1, index, diff)

    def _update(self, node, start, end, index, diff):
        if index < start or index > end:
            return
        self.tree[node] += diff
        if start != end:
            mid = (start + end) // 2
            self._update(2 * node, start, mid, index, diff)
            self._update(2 * node + 1, mid + 1, end, index, diff)

    def query(self, left, right):
        return self._query(1, 0, len(self.arr) - 1, left, right)

    def _query(self, node, start, end, left, right):
        if left > end or right < start:
            return 0
        if left <= start and end <= right:
            return self.tree[node]
        mid = (start + end) // 2
        return self._query(2 * node, start, mid, left, right) + self._query(2 * node + 1, mid + 1, end, left, right)

• init 메서드 :
  • 세그먼트 트리를 초기화
  • arr는 세그먼트 트리를 초기화할 때 사용할 원래 배열, 이 배열의 각 요소는 세그먼트 트리의 각 리프 노드에 대응
  • => 배열의 인덱스와 세그먼트 트리의 노드 번호를 일치시키는 것이 가능
  • 세그먼트 트리는 [0] * (4 * len(arr))와 같이 배열로 구현되며, 세그먼트 트리의 노드 번호는 1부터 시작
• build 메서드 :
  • 세그먼트 트리를 빌드
  • 트리를 구축하는 데 있어서 분할 정복 알고리즘을 사용
  • 트리의 리프 노드에 arr의 각 요소를 할당
    • 현재 노드가 리프 노드인 경우(start == end), 노드에는 해당 위치의 원래 배열(arr)의 값을 저장한다.
  • 트리의 각 노드에 대해서, 그 노드의 두 자식 노드에 대한 정보를 이용하여 그 노드의 값을 계산
    • 현재 노드는 두 자식 노드(2 * node, 2 * node + 1)의 합을 저장
  • 모든 노드에 대해서 반복
• update 메서드 : 
  • 세그먼트 트리에서 특정 위치의 값을 업데이트
  • 세그먼트 트리의 리프 노드부터 시작하여 해당 위치까지 올라가면서, 각 노드의 값을 갱신
    • 배열(arr)에서 index번째 위치의 값을 value로 바꾼 후, _update 메서드를 호출
    • _update 메서드는 세그먼트 트리에서 index번째 위치에 해당하는 노드부터 루트 노드까지의 구간에 대해 값을 갱신
    • 값의 갱신은 diff만큼만 수행하면 되는데, diff는 새로운 값과 기존 값의 차이
  • 이때, 노드의 값을 갱신하면서, 그 노드의 부모 노드에 대한 값을 갱신
• query 메서드 : 
  • 세그먼트 트리에서 특정 구간의 값을 찾아내는 데 사용
  • 분할 정복 알고리즘을 사용
  • 먼저, 찾고자 하는 구간이 트리의 현재 노드가 표현하는 구간과 완전히 일치하면, 그 노드의 값을 반환
    • _query 메서드를 호출하여 이를 구현하는데, 현재 노드(node)가 나타내는 구간과 특정 구간(left부터 right까지)이 겹치지 않는 경우는 0을 반환
    • 현재 노드가 특정 구간(left부터 right까지)에 완전히 포함되는 경우는 현재 노드의 값을 반환
  • 그렇지 않다면, 노드의 두 자식 노드에 대한 정보를 이용하여, 찾고자 하는 구간을 더 작은 구간으로 분할한 다음, 각각의 자식 노드에 대해 query 메서드를 호출
    • 노드의 두 자식 노드(2 * node, 2 * node + 1)에 대해 _query를 호출한 결과를 합한 값을 반환

 

+) query, _query로 나눈이유 :

(앞에 _가 붙은 메서드는 해당 클래스 외부에서 직접 호출하지 않도록 권장되는 private 메서드를 의미한다.)

 

query 메서드를 호출할 때마다 _query를 재귀적으로 호출하며 세그먼트 트리를 탐색하기 때문이다.

 

query 메서드는 외부에서 호출되는 메서드이기 때문에 사용자가 원하는 쿼리 범위를 체크하고, _query를 호출하기 전에 유효성을 검사해야 한다.

이를 위해 _query를 호출하는 _query를 하나 더 만들어서 사용자가 원하는 쿼리 범위가 유효한 경우에만 _query를 호출하도록 구현한 것이다.

이를 통해 메서드 호출 시마다 유효성을 검사하는 것을 방지하고, 코드의 가독성을 높일 수 있다.

 

📖 reference

세그먼트 트리(Segment Tree) 개념 및 구현