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

🦕 세그먼트 트리(Segment Tree)란?

세그먼트 트리는 배열에서 특정 구간에 대한 정보를 빠르게 쿼리하거나 수정할 수 있게 해주는 이진 트리 기반의 자료구조입니다. 이 자료구조는 주로 범위 쿼리 문제에 활용되며, 그 중에서도 구간 합을 빠르게 계산하는 데에 자주 사용됩니다.

원리
세그먼트 트리의 각 노드는 배열의 특정 구간을 대표합니다.
리프 노드는 배열의 개별 원소를 나타냅니다.
부모 노드는 자식 노드들의 합을 저장합니다.
이런 방식으로 트리의 루트 노드는 배열 전체의 합을 저장하게 됩니다.

예시
배열 A [1, 3, 5, 7, 9, 11]을 고려해봅시다.
이때 Tree 는 [0, 36, 9, 27, 4, 5, 16, 11, 1, 3, 7, 9] 가 됩니다.

각 노드에는 배열의 구간 합이 저장되어 있습니다.
리프 노드에는 주어진 배열(A) 값들이 저장되고 내부 노드에는 자식 노드의 합이 저장됩니다.

루트 노드(36)는 전체 배열의 합을 나타냅니다.
좌측 자식 노드(9)는 배열의 첫 번째부터 세 번째 원소까지의 합(A[0:2])을 나타냅니다.
우측 자식 노드(27)는 배열의 네 번째부터 여섯 번째 원소까지의 합(A[3:5])을 나타냅니다.
이와 같이 세그먼트 트리는 구간의 합을 빠르게 계산하거나, 배열의 특정 원소의 값을 변경할 때도 O(logN) 시간복잡도를 유지하면서 구간의 합을 업데이트할 수 있게 해줍니다.

---

👈🏿 파이썬 코드 구현

1. init 메소드

def __init__(self, arr):
    # 배열의 길이
    self.n = len(arr)
    # 세그먼트 트리의 사이즈는 원래 배열의 4배로 설정 (충분한 공간 확보)
    self.tree = [0] * (4 * self.n)
    # 트리 빌드
    self.build(0, self.n-1, 1, arr)

1. self.n = len(arr):
   주어진 배열의 길이를 저장합니다. 이는 나중에 세그먼트 트리의 범위를 정할 때 필요합니다.
2. self.tree = [0] * (4 * self.n):
   세그먼트 트리를 저장할 리스트를 초기화합니다. 세그먼트 트리의 크기는 원래 배열의 길이의 4배가 되도록 설정합니다. 이렇게 설정하는 이유는 세그먼트 트리가 완전 이진 트리의 형태를 가지며, 이를 배열로 표현할 때 충분한 공간을 확보하기 위함입니다.
3. self.build(0, self.n-1, 1, arr):
   주어진 배열을 기반으로 세그먼트 트리를 구축하는 메소드를 호출합니다. 이 때, 전체 배열의 범위인 0부터 self.n-1까지를 대상으로 하며, 세그먼트 트리의 루트 노드는 1이라는 인덱스를 가집니다.

---

2. build 메소드

build 메소드는 초기 배열을 받아 세그먼트 트리를 구축합니다.

분할 정복 알고리즘을 이용합니다.

def build(self, start, end, node, arr):
    # 리프 노드인 경우 배열의 값을 그대로 대입
    if start == end:
        self.tree[node] = arr[start]
        return self.tree[node]
    # 구간을 절반으로 나누어 재귀적으로 트리 생성
    mid = (start + end) // 2
    self.tree[node] = self.build(start, mid, node*2, arr) + self.build(mid+1, end, node*2+1, arr)
    return self.tree[node]

• start, end는 현재 구간을 나타냅니다. 처음 호출 시 전체 배열 범위를 의미합니다.
• node는 현재 세그먼트 트리에서의 위치 (노드 번호)를 의미합니다.
  

동작방식

1. 만약 start와 end가 동일하면 (즉, 리프 노드에 도달했다면) 해당 노드에 배열의 해당 인덱스 값을 저장합니다.
2. 그렇지 않다면, 현재 구간을 두 부분으로 나누고 재귀적으로 왼쪽과 오른쪽 자식 노드를 빌드합니다.
3. 현재 노드의 값을, 왼쪽 자식 노드와 오른쪽 자식 노드의 합으로 설정합니다.

---

3. query 메소드

query 메소드는 지정된 구간의 합을 계산합니다.

def query(self, left, right, node, node_left, node_right):
    # 구간이 겹치지 않는 경우
    if right < node_left or node_right < left:
        return 0
    # 현재 노드 구간이 쿼리 구간에 완전히 포함될 경우
    if left <= node_left and node_right <= right:
        return self.tree[node]
    # 그렇지 않다면 두 부분으로 나누어 합을 계산
    mid = (node_left + node_right) // 2
    return self.query(left, right, node*2, node_left, mid) + self.query(left, right, node*2+1, mid+1, node_right)

query 메소드는 지정된 구간의 합을 계산합니다.

• left, right는 구간 합을 구하고자 하는 범위입니다.
• node_left, node_right는 현재 노드가 커버하는 범위입니다.

동작 방식

1. 만약 현재 노드의 범위와 요청된 범위가 겹치지 않으면, 0을 반환합니다.
2. 요청된 범위가 현재 노드의 범위에 완전히 포함된다면, 현재 노드의 값을 반환합니다.
3. 그렇지 않다면, 요청된 범위를 두 부분으로 나누어 재귀적으로 합을 계산하고, 두 결과를 합쳐 반환합니다.

---

4. update 메소드

update 메소드는 배열의 특정 원소 값을 변경하고, 이에 따라 세그먼트 트리를 갱신합니다.

    def update(self, index, new_value, node, node_left, node_right):
        # 구간 밖인 경우
        if index < node_left or node_right < index:
            return self.tree[node]
        # 리프 노드인 경우 값을 갱신
        if node_left == node_right:
            self.tree[node] = new_value
            return self.tree[node]
        # 리프 노드가 아닌 경우, 자식 노드들도 갱신해야 함
        mid = (node_left + node_right) // 2
        self.tree[node] = self.update(index, new_value, node*2, node_left, mid) + self.update(index, new_value, node*2+1, mid+1, node_right)
        return self.tree[node]

• index는 갱신하려는 원소의 위치입니다.
• new_value는 갱신하려는 값입니다.

동작 방식

1. 만약 현재 노드의 범위가 갱신하려는 원소를 포함하지 않으면, 현재 노드의 값을 반환합니다.
2. 만약 리프 노드에 도달했다면, 원소의 값을 갱신합니다.
3. 그렇지 않다면, 해당 원소를 포함하는 자식 노드들을 재귀적으로 갱신하고, 현재 노드의 값을 자식 노드들의 합으로 갱신합니다.

---

5. 전체 코드

class SegmentTree:
    def __init__(self, arr):
        # 배열의 길이
        self.n = len(arr)
        # 세그먼트 트리의 사이즈는 원래 배열의 4배로 설정 (충분한 공간 확보)
        self.tree = [0] * (4 * self.n)
        # 트리 빌드
        self.build(0, self.n-1, 1, arr)

    # 세그먼트 트리 빌드 함수
    def build(self, start, end, node, arr):
        # 리프 노드인 경우 배열의 값을 그대로 대입
        if start == end:
            self.tree[node] = arr[start]
            return self.tree[node]
        # 구간을 절반으로 나누어 재귀적으로 트리 생성
        mid = (start + end) // 2
        self.tree[node] = self.build(start, mid, node*2, arr) + self.build(mid+1, end, node*2+1, arr)
        return self.tree[node]

    # 구간 합을 계산하는 함수
    def query(self, left, right, node, node_left, node_right):
        # 구간이 겹치지 않는 경우
        if right < node_left or node_right < left:
            return 0
        # 현재 노드 구간이 쿼리 구간에 완전히 포함될 경우
        if left <= node_left and node_right <= right:
            return self.tree[node]
        # 그렇지 않다면 두 부분으로 나누어 합을 계산
        mid = (node_left + node_right) // 2
        return self.query(left, right, node*2, node_left, mid) + self.query(left, right, node*2+1, mid+1, node_right)

    # 배열의 특정 원소를 갱신하는 함수
    def update(self, index, new_value, node, node_left, node_right):
        # 구간 밖인 경우
        if index < node_left or node_right < index:
            return self.tree[node]
        # 리프 노드인 경우 값을 갱신
        if node_left == node_right:
            self.tree[node] = new_value
            return self.tree[node]
        # 리프 노드가 아닌 경우, 자식 노드들도 갱신해야 함
        mid = (node_left + node_right) // 2
        self.tree[node] = self.update(index, new_value, node*2, node_left, mid) + self.update(index, new_value, node*2+1, mid+1, node_right)
        return self.tree[node]

# arr의 값 설정
A = [1, 3, 5, 7, 9, 11]
seg_tree = SegmentTree(A)

# 예시 쿼리 및 출력
print(seg_tree.query(1, 3, 1, 0, seg_tree.n-1))  # 1~3 인덱스 구간 합
# 값 변경 예시: arr의 2번째 원소를 10으로 변경
seg_tree.update(2, 10, 1, 0, seg_tree.n-1)
print(seg_tree.query(1, 3, 1, 0, seg_tree.n-1))  # 1~3 인덱스 구간 합 출력

A = [1, 3, 5, 7, 9, 11]
15 [3, 5, 7]

A = [1, 3, 10, 7, 9, 11]
20 [3, 10, 7]