세그먼트 트리 (Segment Tree)

 세그먼트 트리는 구간 합을 구할 때 사용할 수 있다. 구간 합을 단순 반복문으로 linear하게 계산하면 O(N) 시간이 소요된다. 세그먼트 트리를 사용하면 O(log N) 시간에 구할 수 있어 더 빠르다.

 

출처 : https://www.geeksforgeeks.org/segment-tree-data-structure/

 위 사진은 0부터 5까지의 인덱스를 가진 배열, 즉 6개의 숫자가 들어있는 배열을 세그먼트 트리로 만든 것이다. 이진 탐색(Binary Search)의 메커니즘을 이용하여 트리를 만든다. 세그먼트 트리는 배열을 이용한 트리 구현 방법을 사용하는 것이 편하다. 이 때, 트리의 인덱스를 1부터 시작하게 만들면 index * 2 는 왼쪽 자식 노드를 가리키고, index * 2 + 1 은 오른쪽 자식 노드를 가리키게 되어 간편하다. 참고로 입력받을 배열도 인덱스를 1부터 시작하게 만드는 것이 구현하기 더 편하다.

 위 사진을 보면 입력 배열의 크기는 6인데 반해 세그먼트 트리의 노드 수는 11개이다. 세그먼트 트리의 크기는 입력 배열의 크기의 4배로 만들면 된다.

 루트 노드(index 1)에는 처음부터 끝까지의 합을 나타내고, 왼쪽 자식 노드는 처음부터 중간까지, 오른쪽 자식 노드는 (중간 + 1)부터 끝까지의 합을 나타낸다. 이진 탐색이 보통 재귀를 이용하여 수행되는 것처럼, 세그먼트 트리를 만들 때도 재귀를 사용한다.

int64_t makeSegmentTree(int start, int end, int node) {
	int mid;
	if (start == end) {	// 리프 노드 도착
		return segmentTree[node] = numbers[start];
	}
	mid = (start + end) / 2;
	return segmentTree[node] = makeSegmentTree(start, mid, node * 2) 
		+ makeSegmentTree(mid + 1, end, node * 2 + 1);
}

numbers 는 처음 입력받았던 배열이다. 

 

 세그먼트 트리에서 구간 합을 구하는 것은 각 노드가 가리키는 구간(start, end)이 구하고자 하는 구간(left, right) 안에 들어가 있는 모든 노드를 더하면 된다.

int64_t findSum(int start, int end, int node, int left, int right) {
	int mid;
	if (right < start || left > end) {	// 범위 밖
		return 0;
	}
	if (left <= start && right >= end) {	// 범위 안
		return segmentTree[node];
	}
	mid = (start + end) / 2;
	return findSum(start, mid, node * 2, left, right) 
		+ findSum(mid + 1, end, node * 2 + 1, left, right);
}

첫번째 if 문은 노드의 구간이 아예 벗어나 있을 경우 구할 필요가 없으므로 바로 0을 리턴하고, 두번째 if 문은 노드의 구간이 구하고자 하는 구간에 완전히 포함되어 있는 경우이고 이 때는 그냥 노드의 값을 리턴하면 된다. 마지막 세 줄은 걸쳐있는, 즉 약간 포함되어 있는 경우이므로 자식 노드 둘을 체크하게 된다. 

 

 처음 입력받았던 배열의 한 요소가 수정되었을 경우, 세그먼트 트리도 역시 수정해주어야 한다. 이 때는 수정된 요소의 인덱스가 포함되어있는 모든 구간을 수정해주면 된다. 

void update(int start, int end, int node, int index, int64_t num) {
	int mid;
	if (index < start || index > end) {	// 범위 밖
		return;
	}
	segmentTree[node] += num;
	if (start == end) {	// base case
		return;
	}
	mid = (start + end) / 2;
	update(start, mid, node * 2, index, num);
	update(mid + 1, end, node * 2 + 1, index, num);
}

여기서 주의해야 할 점은 num이 바뀐 숫자가 아니라는 점이다. segmentTree[node] += num 에서 알 수 있듯이 바뀐 만큼의 차이를 더하게 된다. 처음 update 함수의 num 매개변수를 넣어줄 때, (수정된 값 - 원래의 값)을 넣어주면 수정된 값이 더 클 경우에는 num이 양수가 되고, 수정된 값이 원래의 값보다 더 작을 때는 음수가 들어가게 되므로 그 차이만큼 각 노드의 구간 합도 올바르게 수정된다.