세그먼트 트리

세그먼트 트리

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

더 큰 범위로는 '인덱스 트리'라고 불리는데, 코딩 테스트 영역에서는 큰 차이가 없다.

 

세그먼트 트리 핵심 이론

세그먼트 트리의 종류는 구간 합, 최대/최소 구하기로 나눌 수 있고다.
구현 단계는 트리 초기화하기 -> 질의값 구하기(구간 합 or 최대/최소) -> 데이터 업데이트하기

 

 

세그먼트 트리는 이론이 길었다.. 휴

그래도 쓰고, 그림그려보면서 이해해보니 주어진 데이터를 세그먼트 트리로 표현하고, 인덱스 조작을 잘 하면서 노드 값을 컨트롤하면 원하는 구간 합이나, 구간의 최대 최소를 빠른 시간 복잡도로 구할 수 있는 알고리즘이다.

 

문제를 풀어보면서 더 이해해보자.

 

백준 2042번 : 구간 합 구하기

https://www.acmicpc.net/problem/2042

#include <iostream>
#include <vector>
using namespace std;

using INT64 = long long;
int N, M, K;
vector<INT64> tree;

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	
	cin >> N >> M >> K;

	//1. 트리 배열 초기화
	//트리 배열의 크기 구하기
	int k = 1;
	while ((1 << k) < N) 
		k++; // 높이 계산
	tree.resize(pow(2, k) * 2, 0);

	int leafNodeStartIdx = pow(2, k); //리프노드가 시작되는 인덱스
	//리프 노드 채우기
	for (int i = 0; i < N; i++)
	{
		cin >> tree[leafNodeStartIdx + i];
	}
	
	//구간 합으로 나머지 배열도 채운다.
	for (int i = leafNodeStartIdx - 1; i > 0; i--)
	{
		tree[i] = tree[i * 2] + tree[i * 2 + 1];
	}

	for (int i = 0; i < M + K; i++)
	{
		INT64 a, b, c;
		cin >> a >> b >> c;
		
		if (a == 1) //업데이트(b번째 수를 c로 바꿔라)
		{
			//주어진 질의 인덱스를 트리의 인덱스로 변환
			int targetIdx = b + pow(2, k) - 1;
			
			//부모로 타고 올라가면서 변화 반영
			tree[targetIdx] = c; // 변경된 값 설정
			while (targetIdx > 1)
			{
				targetIdx /= 2;
				tree[targetIdx] = tree[targetIdx * 2] + tree[targetIdx * 2 + 1];
			}
		}
		else //구간 합 구하기(구간 b~c의 합을 구해라)
		{
			int start_idx = b + pow(2, k) - 1;
			int end_idx = c + pow(2, k) - 1;
			INT64 sum = 0;
			while (start_idx <= end_idx)
			{
				//선택 여부 판단
				if (start_idx % 2 == 1)
					sum += tree[start_idx];
				if (end_idx % 2 == 0)
					sum += tree[end_idx];

				//depth 변경
				start_idx = (start_idx + 1) / 2;
				end_idx = (end_idx - 1) / 2;
			}
			cout << sum << '\n';
		}
	}

	return 0;
}

 

이 문제는 노드 개수가 충분히 크기 때문에 단순한 방법으로 업데이트와 구간 합을 구하면 시간 복잡도가 터진다.

심지어 중간에 수 변경이 빈번하게 변경되기도 하므로 합 배열로 구간합을 구하는 것은 어렵다.

내용이 빈번하게 변하고 구간 합을 구하는 문제이므로 세그먼트 트리로 해결 할 수 있다.

 

위에서 배운 이론대로 구현을 쭉 하면 되는데

구현 상에 실수했던 점은 k를 구할 때 log2() 함수로 구했는데 그 방법보다 저렇게 직관적으로 << k 연산으로 N보다 커질 때의 k 값을 구하면 되었다.

그리고 부모 노드에 업데이트 사항을 전파할 때도 복잡하게 하지 말고 그냥 변경될 값으로 덮어 씌어버린 후 부모타고 올라가면서 자식 노드 두 개씩 더하는 방식을 활용하면 직관적으로 구현할 수 있다.

 

백준 10868번 : 최솟값

https://www.acmicpc.net/problem/10868

#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

using INT64 = long long;
int N, M;
vector<INT64> tree;

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	
	cin >> N >> M;

	//1. 트리 배열 초기화
	//트리 배열의 크기 구하기
	int k = 1;
	while ((1 << k) < N) 
		k++; // 높이 계산
	tree.resize(pow(2, k) * 2, 0);

	int leafNodeStartIdx = pow(2, k); //리프노드가 시작되는 인덱스
	//리프 노드 채우기
	for (int i = 0; i < N; i++)
	{
		cin >> tree[leafNodeStartIdx + i];
	}
	
	//최소 값으로 나머지 배열도 채운다.
	for (int i = leafNodeStartIdx - 1; i > 0; i--)
	{
		tree[i] = min(tree[i * 2], tree[i * 2 + 1]);
	}

	for (int i = 0; i < M; i++)
	{
		INT64 a, b;
		cin >> a >> b;
		
		int start_idx = a + pow(2, k) - 1;
		int end_idx = b + pow(2, k) - 1;

		INT64 minValue = 1000000000;
		while (start_idx <= end_idx)
		{
			//선택 여부 판단
			if (start_idx % 2 == 1)
				minValue = min(minValue, tree[start_idx]);
			if (end_idx % 2 == 0)
				minValue = min(minValue, tree[end_idx]);

			//depth 변경
			start_idx = (start_idx + 1) / 2;
			end_idx = (end_idx - 1) / 2;
		}
		cout << minValue << '\n';
	}
	return 0;
}

 

이 문제는 구간의 최소 값을 묻기 때문에 세그먼트 트리로 풀 수 있는 문제이다.

이론대로 구현하면 풀 수 있다.

세그먼트 트리는 구현 자체가 조금 복잡하다보니 난이도가 높다.

 

백준 11505번 : 구간 곱 구하기

https://www.acmicpc.net/problem/11505

#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

using INT64 = long long;
int N, M, K;
vector<INT64> tree;
constexpr INT64 MOD_VAL = 1000000007;

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	
	cin >> N >> M >> K;

	//1. 트리 배열 초기화
	//트리 배열의 크기 구하기
	int k = 1;
	while ((1 << k) < N) 
		k++; // 높이 계산
	tree.resize(pow(2, k) * 2, 0);

	int leafNodeStartIdx = pow(2, k); //리프노드가 시작되는 인덱스
	//리프 노드 채우기
	for (int i = 0; i < N; i++)
	{
		cin >> tree[leafNodeStartIdx + i];
	}
	
	//구간 곱으로 나머지 배열도 채운다.
	for (int i = leafNodeStartIdx - 1; i > 0; i--)
	{
		tree[i] = tree[i * 2] * tree[i * 2 + 1] % MOD_VAL;
	}

	for (int i = 0; i < M + K; i++)
	{
		INT64 a, b, c;
		cin >> a >> b >> c;
		
		if (a == 1) // 업데이트
		{
			int target_idx = b + pow(2, k) - 1;
			tree[target_idx] = c;
			while (target_idx > 0)
			{
				target_idx /= 2;
				tree[target_idx] = tree[target_idx * 2] * tree[target_idx * 2 + 1] % MOD_VAL;
			}
		}
		else // 구간 곱
		{
			int start_idx = b + pow(2, k) - 1;
			int end_idx = c + pow(2, k) - 1;
			INT64 total = 1;
			while (start_idx <= end_idx)
			{
				//선택 여부 판단
				if (start_idx % 2 == 1)
					total = total * tree[start_idx] % MOD_VAL;
				if (end_idx % 2 == 0)
					total = total * tree[end_idx] % MOD_VAL;

				//depth 변경
				start_idx = (start_idx + 1) / 2;
				end_idx = (end_idx - 1) / 2;
			}
			cout << total << '\n';
		}
	}
	return 0;
}

 

이 문제는 구간 곱이다. 위에서 풀었던 세그먼트 트리 문제들을 약간 변형하면 풀 수 있는 문제다.

이론을 잘 이해했다면, 쉽게 구현할 수 있다.

 

공부 자료 : Do it! 코딩 테스트 알고리즘 C++(김종관, 이지스 퍼블리싱)