세그먼트 트리

세그먼트 트리

세그먼트 트리는 Range Query (Update, Sum, Max, Etc…)를 처리하는데 효율적인 자료구조로, 완전 이진 트리이며 인덱스 트리이기 때문에 1차원 배열으로 표현되는 트리입니다. 인덱스는 1부터 시작하며 자식 노드는 2 n, 2 n + 1 번 인덱스를 가지고, 부모 노드는 n // 2 입니다. 세그먼트 트리 빌드에 $O(n)$, 쿼리에 $O(\log n)$의 시간복잡도를 가지기 때문에 효율적으로 데이터를 관리할 수 있습니다.

세그먼트 트리는 원래의 데이터를 리프 노드로 가지고, 그 노드를 두개씩 Merge 하면서 부모노드에 저장합니다. 이때 Merge를 하여 부모노드에 저장하기 때문에 1번부터 4번까지의 쿼리를 구한다고 하면, 미리 Merge 되어있는 부모노드를 쿼리로 반환하면 되기 때문에 효율적이라고 할 수 있습니다.

Build

세그먼트 트리는 원래의 데이터들을 리프노드로 가지고, 두개씩 Merge해 부모노드에 저장합니다.

def merge(left, right):
	return sum(left, right) # max(left, right), gcd(left, right) ...

default = 0 # default
segment_tree = [default] * length * 4 # 세그먼트 트리의 크기는 데이터의 개수 * 4
def build(arr, node, start, end):
	# arr : 데이터 배열
	# node : 몇번째 노드
	# start, end : 데이터 배열에서의 start, end 범위
	
	if start == end: # 만약 범위가 1이라면 리프노드이므로, 세그먼트 트리의 해당 노드에 해당 데이터를 저장
		segment_tree[node] = arr[start]
		return segment_tree[node] # 리프 노드를 반환

	mid = (start + end) // 2 # 세그먼트 트리는 이진 트리이기 때문에 왼쪽 반과 오른쪽 반을 나눔
	segment_tree[node] = merge(build(arr, node * 2, start, mid), build(arr, node * 2 + 1, mid + 1, end))
	return segment_tree[node]
	# 만약 리프노드가 아니라면 자식 노드 두개를 Merge 한 것을 부모노드에 저장

Range Query

부모노드부터 그 범위에 포함되는지 확인 후 완전히 포함된다면 그것을 반환하게 합니다.

def query(left, right, node, start, end):
	# left, right : 쿼리의 범위
	if end < left or right < start: # start가 right보다 오른쪽이거나 end가 left 보다 왼쪽이라면 범위를 완전히 벗어난 것이므로 default 값을 반환
		return default
	if start <= left and right <= end: # 범위에 완전히 포함될 경우 그 노드 값을 반환
		return segment_tree[node]
	
	mid = (start + end) // 2
	return merge(query(left, right, node * 2, start, mid), query(left, right, node * 2 + 1, mid + 1, end))

Range Update

업데이트 또한 Range Query 와 같이 범위에 따라서 업데이트 합니다. 그러나, 업데이트 할 때 하나의 원소만 업데이트 하기 때문에 리프노드 까지 갔을 때 그 값을 변경해주어야합니다.

def update(index, value, node, start, end):
	# left, right : 쿼리의 범위
	if index < start or index > end: # start가 right보다 오른쪽이거나 end가 left 보다 왼쪽이라면 범위를 완전히 벗어난 것이므로 default 값을 반환
		return segment_tree[node] # 업데이트가 일어나지 않기 때문에 그대로의 값을 반환
	if start == end:
		segment_tree[node] = value
		return segment_tree[node]
	
	mid = (start + end) // 2
	segment_tree[node] = merge(update(index, value, node * 2, start, mid), query(index, value, node * 2 + 1, mid + 1, end))
	return segment_tree[node]

이때 구간을 변경하는 경우 Lazy Propagation을 사용해야합니다.