세그먼트 트리(Segment Tree, 구간 트리)

(출처: 종만북)

세그먼트 트리(Segment Tree)는 저장된 자료를 전처리해서, 특정 구간에 대한 쿼리를 빠르게 수행할 수 있는 자료구조입니다.

 

배열에 N개의 원소가 있다고 했을 때 구현 방식에 따라 2N에서 4N까지 메모리가 필요하지만, 원소의 변경이나 특정 범위 내 원소의 연산을 O(lgN)에 수행할 수 있는 장점이 있습니다.

세그먼트 트리는 수행하는 쿼리에 따라 여러가지 구현방식이 존재하는데, 이번 게시물에서는 1차원 배열에서 특정 구간의 최소치를 찾는 쿼리인 구간 최소 쿼리(Range Minimum Query, RMQ)를 예로 들겠습니다.

 

 

Fig 1. 길이 15인 배열을 표현하는 세그먼트 트리의 노드가 저장하는 구간들

세그먼트 트리의 기본 아이디어는, 주어진 배열의 구간을 표현하는 이진 트리를 만드는 것입니다.

Fig 1.은 배열의 길이가 15일 때 세그먼트 트리의 각 노드가 표현하는 구간을 나타낸 것입니다. 맨 위 노드가 루트이고, 아래는 자식 노드들을 나타냅니다.

 

세그먼트 트리는 각 구간에 대한 계산 결과를 노드에 저장합니다. 예를 들면, 구간의 최소치를 구하는 세그먼트 트리는 해당되는 구간의 최소값을 각 노드에 저장합니다.

이 과정을 거쳐 어떤 구간 \([i, j]\)가 주어졌을 때 이를 세그먼트 트리 내부 구간들의 합집합으로 나타낼 수 있습니다. Fig 1.은 \([2,8]\) 구간을 각각 \([2,3]\), \([4,7]\), \([8,8]\)의 합집합으로 나타낸 것입니다. 구간에 포함되는 노드를 노란색 배경으로 나타냈습니다. 각 구간의 최소값은 미리 계산되어있기 때문에 각 구간의 값 중 최소값을 고르면 쿼리의 답이 됩니다.

 

 

세그먼트 트리 구현

구현 방식은 Top-down 방식과 Bottom-up 방식이 존재하는데, Top-Down 방식을 기준으로 설명하겠습니다.

 

Fig 2. N=6인 배열에 대한 세그먼트 트리

먼저 루트 노드를 1번 원소로 하고 i번 노드의 왼쪽 자식과 오른쪽 자식을 2*i, 2*i+1번 노드로 나타낼 수 있습니다. 이렇게 하면 힙과 비슷하게 배열로 만들어 표현할 수 있습니다.

배열의 크기가 2의 제곱꼴이면 Perfect Binary Tree이기 때문에 이 때 트리의 높이는 H=lgN이고, leaf 노드가 N개인 경우 필요한 배열의 크기는 2N입니다(0번 노드를 비우지 않으면 2N-1입니다). 배열의 길이가 2의 제곱꼴이 아니면 트리의 높이는 \(H=\left\lceil \log_2{N} \right\rceil \)이고, 이때 필요한 배열의 크기는 \(2^{H+1}\) 입니다.

계산이 귀찮으면 4N이 모든 경우에 필요한 배열의 크기보다 크기 때문에 메모리의 낭비가 좀 있지만 4N을 사용해도 무방합니다.

 

Fig 2.는 N=6인 배열에 대한 세그먼트 트리 예시입니다. H=3이기 때문에 세그먼트 트리를 만들기 위해 2^4=16 크기의 배열이 필요합니다. 10, 11번 노드가 비어있는 점을 참고해주세요.

 

 

const int INT_MAX = numeric_limits<int>::max();

// range minimum query
struct RMQ {
private:
	int n;
	vector<int> rangeMin;
	
	// node의 표현 범위 arr[nodeLeft, nodeRight]가 주어질 때
	// 이 범위와 arr[left, right] 교집합의 최소치를 구함
	int query(int left, int right,
		int node, int nodeLeft, int nodeRight) {
		// 두 구간이 겹치지 않은 경우
		if (right < nodeLeft || nodeRight < left) return INT_MAX;

		// node 표현 범위가 arr[left, right]에 포함되는 경우
		if (left <= nodeLeft && nodeRight <= right) return rangeMin[node];

		int mid = (nodeLeft + nodeRight) / 2;
		return min(query(left, right, node*2, nodeLeft, mid),
			query(left, right, node*2+1, mid+1, nodeRight));
	}

	// array[index] = newValue로 바뀌었을 때 node를 루트로 하는 구간트리를 갱신
	// 노드가 표현하는 구간의 최소치를 반환
	int update(int index, int newValue,
		int node, int nodeLeft, int nodeRight) {
		// index가 노드가 표현하는 구간과 상관 없는 경우
		if (index < nodeLeft || nodeRight < index) return rangeMin[node];

		// 트리의 리프까지 내려온 경우
		if (nodeLeft == nodeRight) return rangeMin[node] = newValue;

		int mid = (nodeLeft + nodeRight) / 2;

		return rangeMin[node] = min(update(index, newValue, node * 2, nodeLeft, mid),
			update(index, newValue, node * 2 + 1, mid + 1, nodeRight));
	}

public:
	RMQ(const vector<int>& array) {
		n = array.size();
		rangeMin.resize(n * 4);
		init(array, 0, n - 1, 1);
	}

	// node가 array[left..right] 배열을 표현할 때
	// node를 루트로 하는 서브트리를 초기화하고, 이 구간의 최소치 반환
	int init(const vector<int>& array, int left, int right, int node) {
		if (left == right)
			return rangeMin[node] = array[left];

		int mid = (left + right) / 2;
		int leftMin = init(array, left, mid, node * 2);
		int rightMin = init(array, mid+1, right, node * 2 + 1);

		return rangeMin[node] = min(leftMin, rightMin);
	}

	int query(int left, int right) {
		return query(left, right, 1, 0, n - 1);
	}
	
	int update(int index, int newValue) {
		return update(index, newValue, 1, 0, n - 1);
	}
};

구현 설명은 다음과 같습니다.

1. 초기화 (init): 배열을 입력받아 4*n 크기의 배열을 할당합니다. 배열의 1번 원소를 root로 사용하고, 재귀적으로 자식 노드를 초기화합니다. init()은 현재 구간을 두 개로 나눠 재귀 호출한 뒤 두 구간의 최소치 중 더 작은 값을 선택해 현재 노드 구간의 최소값을 계산합니다.

2. 질의 처리 (query): node가 표현하는 구간 [nodeLeft, nodeRight]와 최소치를 찾을 구간 [left, right]의 교집합을 구해 그에 따라 서로 다른 값을 반환합니다.

• 교집합이 공집합인 경우: 두 구간이 서로 겹치지 않으니 반환 값이 무시되는 값을 반환해줍니다.
  (RMQ에서는 아주 큰값 반환)
• 교집합이 [nodeLeft, nodeRight]인 경우: [left, right]가 노드 표현 범위를 완전히 포함하니 미리 계산된 최소값을 반환하면 됩니다
• 그 외의 경우: 두 자식 노드에 대해 query를 재귀호출한 뒤 두 값중 작은 값을 반환합니다.

3. 갱신 (update): 세그먼트 트리를 생성한 뒤 배열의 값이 바뀐 경우, 세그먼트 트리를 다시 만들 수도 있지만, 값이 하나만 바뀐 경우 세그먼트 트리를 빠르게 갱신할 수 있습니다.

갱신 과정은 초기화와 질의 처리 메소드를 합친 것처럼 구현됩니다. 해당 노드가 표현하는 구간에 index가 표함되지 않으면 무시하고, 포함되면 자식 노드에 update 메소드를 재귀 호출해서 각 구간의 최소치를 구한 뒤 업데이트합니다.

 

 

참고자료

• 구종만, 프로그래밍 대회에서 배우는 알고리즘 문제해결 전략, 인사이트
• blisstoner, Persistent Segment Tree, Samsung SW Membership Blog
• baekjoon, 세그먼트 트리 (Segment Tree), BOJ