[4주차] 세그먼트 트리

트리가 어떤 것인지 다시 확인하기

https://ysh0129.tistory.com/106

[3주차] 비선형 자료구조

 

목차

Ⅰ. 세그먼트 트리?

Ⅱ. 세그먼트 트리 핵심

1. 트리 초기화 방법
2. 질의값 구하는 방법
3. 데이터 업데이트 방법

 

 

 

Ⅰ. 세그먼트 트리?

▶ 주어진 데이터의 구간 합과 데이터 업데이트를 빠르게 수행하기 위해 고안해낸 자료구조의 형태

(≒ 인덱스 트리)

 

세그먼트 트리로 구간합 인덱스 2~4 구간합 구하는 원리

 

Ⅱ. 세그먼트 트리 핵심

1. 트리 초기화 구조

 

● 크기 설정 : 원본 배열의 크기가 N일 때, 세그먼트 트리를 저장할 배열(리스트)은 보통 4N의 크기로 준비

→  트리가 완전 이진 트리 형태를 유지하도록 충분한 공간을 확보

 

●  초기화 함수 정의 : init , build

매개변수	설명
a	원본 데이터 배열
tree	세그먼트 트리를 저장할 배열
node	현재 트리의 노드 번호(index)
start, end	현재 노드가 담당하는 배열 a의 구간 인덱스

 

# 예시 배열
arr = [1, 2, 3, 4, 5]

# 세그먼트 트리 배열 (크기 4 * n)
n = len(arr)
tree = [0] * (4 * n)

# 세그먼트 트리 초기화 함수
def init(node, start, end):
    # 리프 노드인 경우
    if start == end:
        tree[node] = arr[start]
        return tree[node]
    
    mid = (start + end) // 2
    # 왼쪽, 오른쪽 자식 트리 초기화
    left_sum = init(node * 2, start, mid)
    right_sum = init(node * 2 + 1, mid + 1, end)
    
    # 현재 노드 값 = 왼쪽 + 오른쪽
    tree[node] = left_sum + right_sum
    return tree[node]

# 트리 초기화 실행
init(1, 0, n - 1)

# 결과 출력
print(tree)

 

구분	init	build
기능	초기 데이터로 트리의 모든 노드 값을 채움
주로 사용되는 문맥	클래스의 초기화 과정을 나타낼 때 (내부 함수로 사용)	자료구조를 구축하는 행위 자체를 나타낼 때 (독립 함수로도 사용)
결론	세그먼트 트리를 처음 만들 때 배열 값을 바탕으로 트리를 채우는 동일한 재귀적 로직을 가진다. 이름만 다를 뿐, 기능상 큰 차이는 없다.

 

#init 함수

class SegmentTree:
    def __init__(self, arr):
        # ... (트리 크기 설정)
        self._init(0, n - 1, 1) # 트리를 초기화(구성)하는 내부 함수 호출

    def _init(self, start, end, node):
        # 재귀적으로 트리의 각 노드(구간 합, 최솟값 등)를 채우는 로직
        pass
        
 
#build 함수

	def build_segment_tree(arr):
    	# ... (트리 크기 설정)
    	tree = [0] * tree_size
    	_build(arr, tree, 0, n - 1, 1) # 트리를 구축하는 내부 함수 호출
    	return tree

	def _build(arr, tree, start, end, node):
    	# 재귀적으로 트리의 각 노드(구간 합, 최솟값 등)를 채우는 로직
    	pass

 

 

2. 질의값 구하는 방법

 

구간 질의는 주어진 범위 [L, R] 에 해당하는 값을 트리를 탐색하여 O(\log N) 시간 안에 계산

●  질의값(조회) 함수 정의 : query 

 

매개변수	설명
L (Left Boundary)	사용자가 합을 구하고자 하는 구간의 시작 인덱스
R (Right Boundary)	사용자가 합을 구하고자 하는 구간의 끝 인덱스
start, end	현재 노드가 담당하는 각 시작과 끝 구간 인덱스
node	현재 탐색 중인 세그먼트 트리 배열의 인덱스 (보통 루트 노드에서 1부터 시작)

 

import math

class SegmentTree:
    def __init__(self, arr):
        self.n = len(arr)
        # 트리의 크기를 넉넉하게 4N으로 설정
        tree_size = 1 << (math.ceil(math.log2(self.n)) + 1)
        self.tree = [0] * tree_size
        self.arr = arr
        self._init(0, self.n - 1, 1) # 배열 인덱스: 0 ~ n-1, 트리의 루트 인덱스: 1

    # 트리 구축: arr[start:end+1] 구간의 합을 tree[node]에 저장
    def _init(self, start, end, node):
        if start == end: # 리프 노드
            self.tree[node] = self.arr[start]
            return self.tree[node]
        
        mid = (start + end) // 2
        # 왼쪽 자식 (2*node)과 오른쪽 자식 (2*node + 1)의 합을 현재 노드에 저장
        self.tree[node] = self._init(start, mid, node * 2) + self._init(mid + 1, end, node * 2 + 1)
        return self.tree[node]

    # 구간 합 질의: arr[L:R+1]의 합을 구함
    # L, R: 구하고자 하는 구간의 인덱스 (0-based)
    # start, end: 현재 노드가 담당하는 구간의 인덱스
    # node: 현재 노드의 인덱스
    def query(self, L, R, start, end, node):
        # 1. 완전 불일치: 쿼리 범위가 현재 노드 구간 밖에 있는 경우
        if R < start or L > end:
            return 0
        
        # 2. 완전 포함: 쿼리 범위가 현재 노드 구간을 완전히 포함하는 경우
        if L <= start and end <= R:
            return self.tree[node]
        
        # 3. 부분 포함: 왼쪽 자식과 오른쪽 자식으로 재귀 호출하여 합산
        mid = (start + end) // 2
        left_sum = self.query(L, R, start, mid, node * 2)
        right_sum = self.query(L, R, mid + 1, end, node * 2 + 1)
        
        return left_sum + right_sum

# --- 예시 실행 ---
A = [1, 2, 3, 4, 5, 6, 7, 8]
st = SegmentTree(A)

# 쿼리 예시: 인덱스 2부터 5까지 (A[2]~A[5])의 합 (3+4+5+6 = 18)
# L=2, R=5, 전체 구간은 start=0, end=7, 루트 노드는 node=1
result = st.query(2, 5, 0, st.n - 1, 1)
print(f"배열 {A}에서 인덱스 2부터 5까지의 구간 합: {result}")

# 쿼리 예시 2: 인덱스 0부터 7까지 (전체)의 합 (1+..+8 = 36)
result_all = st.query(0, st.n - 1, 0, st.n - 1, 1)
print(f"전체 구간 합: {result_all}")

 

3. 데이터를 갱신하는 방법

 

배열의 한 점(point)을 변경하고, 그 변경 사항이 영향을 미치는 모든 상위 노드의 값을 수정하는 과정

●  갱신(업데이트) 함수 정의 : update 

매개변수	설명
idx	배열에서 값이 변경된 원본 인덱스
diff	변경된 값의 차이 (new_value - old_value)
start, end	현재 노드가 담당하는 구간
node	현재 트리 노드의 인덱스

 

def update(node, start, end, idx, diff):
    # 범위 밖
    if idx < start or idx > end:
        return
    # 트리 갱신
    tree[node] += diff
    if start != end:
        mid = (start + end) // 2
        update(node * 2, start, mid, idx, diff)
        update(node * 2 + 1, mid + 1, end, idx, diff)

# 예시: arr[2] (3 → 5 로 변경)
idx = 2
diff = 5 - arr[idx]
arr[idx] = 5
update(1, 0, n - 1, idx, diff)

print(query(1, 0, n - 1, 0, 4))  # 전체 합 17로 바뀜