세그먼트 트리 정리

세그먼트 트리는 그냥 일반 알고리즘 수업만 들은 사람이랑 PS를 한다는 사람을 구분짓는 가장 대표적인 알고리즘이 아닐까 싶다. (CLRS 책에 안나와서 그러려나)

인터넷에 쳐보면 그야말로 넘쳐날 정도의 세그먼트 트리 글이 있지만, 아주 쉬운 구현을 발견해서 이를 포함해서 정리 겸 써본다.

세그먼트 트리는 기본적으로 i) 구간에 대한 정보 쿼리, ii) 구간을 한번에 업데이트 시키기 두 가지에 아주 적합한 자료구조이다.

세부적으로 약간 차이가 있지만, 크게 3가지 타입이 있을 것 같다.

i) 단일 업데이트, 구간 쿼리

업데이트가 일어날 때는 하나의 원소만 변경되고, 쿼리는 구간에 대해서 묻는 질의이다.

대표적으로 https://www.acmicpc.net/problem/2042과 같은 문제가 있다.

대표적인 세그먼트 트리 사용 예시이다.

n개의 원소에 대해 세그먼트 트리를 사용할 때는 원소의 개수를 n 이상의 가장 가까운 2의 거듭제곱으로 맞추어 준 뒤 인덱스 1부터 시작하는 것이 구현하기에 편리한데, 그 이유는 p가 "n 이상의 가장 가까운 2의 거듭제곱"이라고 했을 때, tree[p + i] = arr[i]가 성립하기 때문이다. 리프노드의 인덱스를 정확히 알면 bottom-up 방식으로 업데이트나 쿼리를 처리할 수 있게 되는데, 코드가 매우 간결해진다.

그 외 동일한 레벨의 노드는 같은 개수의 원소만을 가진다는 등의 부가적인 장점이 있다.

이렇게 세그먼트 트리를 구현하면 다음과 같다. 구간합을 처리하는 세그먼트 트리인데 내가 본 구현 중 제일 간단하다.

길이 n의 arr에 원소들이 들어있다고 가정한 상태다.

vector<long long> tree;
int p = 1;

void upd(int i, int val) {
	i += p;
	tree[i] = val;
	for (i /= 2; i >= 1; i /= 2) {
		tree[i] = tree[2 * i] + tree[2 * i + 1];
	}
}

long long get(int l, int r) {
	long long ret = 0;
	l += p; r += p;
	while (l <= r) {
		if (l % 2 == 1) ret += tree[l++];
		if (r % 2 == 0) ret += tree[r--];
		l /= 2; r /= 2;
	}
	return ret;
}

while (p < n) p *= 2;
tree.resize(2 * p, 0);

for (int i = 0; i < n; ++i) tree[p + i] = arr[i];
for (int i = p - 1; i >= 1; --i) tree[i] = tree[2 * i] + tree[2 * i + 1];

알고리즘 트레이닝(2019) 책에서 다량의 코드를 참고했다.

세그먼트 트리는 구간에 대한 다양한 쿼리들, 합이나 최대, 최소, and, xor 등을 처리할 수 있다. 그 중 합과 xor은 다음과 같은 성질을 가진다.

query(a, b) = query(0, b) - query(0, a - 1)

이 경우 펜윅 트리라는 것을 사용할 수 있는데, 구현이 세그먼트 트리보다도 훨씬 간단하다.

void insert(int i, int d) {
	while (i <= n) {
		tree[i] += d;
		i += (i & -i);
	}
}

long long get(int i) {
	long long ret = 0;
	while (i) {
		ret += tree[i];
		i &= (i - 1);
	}
	return ret;
}

펜윅 트리는 비트연산을 아름답게 이용한다. 예컨대 i & -i같은 코드는 마지막 비트만 남겨주는 코드이며, i &= (i - 1)은 마지막 비트를 제거하는 코드이다.

펜윅트리의 단점이라면 처음 build를 할 때 각 원소를 모두 insert해주어야 하기 때문에 $O(n\log{n})$의 시간이 든다는 점일 것 같다. 또한, 아쉽게도 max, min, and 등에 대해서는 펜윅트리를 사용할 수 없다.

펜윅트리는 인덱스를 1부터 시작해야 한다는 것에 유의하자.

ii) 구간 업데이트, 단일 쿼리 (합, 곱, XOR만)

펜윅트리를 사용할 수 있는 합과 xor 등은 단일 원소 쿼리일 경우 구간 업데이트 또한 아주 쉽게 할 수 있다.

예를 들어서 1, 2, 3, 5, 8, 1이라는 배열이 있다고 하자. 이 배열에 a~b 사이의 수에 특정 값을 더하고, c번째 원소의 값을 묻는 쿼리 두 종류가 들어온다고 가정해보자.

펜윅트리를 각 원소의 차이인 1, 1, 1, 2, 3, -7로 만들자.

이제 a~b 사이의 수에 d를 더하려면, a번째 원소에 d를 더하고 b+1번째 원소에 -d를 더해주면 간단하게 끝난다.

합, xor 등이 아닌 다른 성질의 것일 경우 어쩔 수 없이 iii)번 형태를 사용해야 한다.

iii) 구간 업데이트, 구간 쿼리

가장 고난도인 형태이다. lazy propagation을 이용해야 한다.

lazy propagation은 아무리 찾아봐도 top-down approach가 구현하기에 더 편한 것 같다.

쿼리의 종류가 엄청 많기 때문에 정형화된 형태가 딱히 없다. 가령 쿼리가 구간의 원소들을 d로 변경하는 게 나올 수도 있고, d만큼 더하는 게 나올 수도 있고, 그럴 때마다 코드를 약간씩 수정해주어야 하기 때문이다. max쿼리 같은 경우 업데이트를 할 때 한 번 내려갔다 올라가며 값을 수정해주어야 하는데 구현이 꽤 까다롭다.

lazy propagation이 성립하려면, 핵심은 lazy 배열과 tree 배열 두 가지를 이용해 해당 노드의 값을 구할 수 있어야 한다. 적용하기 쉽지 않은 예로 a~b 구간의 원소들에 d를 더하는 쿼리와 a~b 구간 원소들의 xor 값을 구하는 쿼리가 같이 들어오는 경우가 있지 않을까 생각한다.

구간의 min값을 묻는 쿼리와, 구간의 원소에 d만큼 더하는 쿼리가 들어왔을 때 대충 잘 작동하는 lazy propagation 코드이다.

class LazyTree {
	private:
		int n;
		vector<long long> tree;
		vector<long long> lazy;
		int build(vector<int>& arr, int i, int li, int ri) {
			if (li == ri) return tree[i] = arr[li];
			int mid = (li + ri) / 2;
			return tree[i] = min(build(arr, 2 * i + 1, li, mid), build(arr, 2 * i + 2, mid + 1, ri));
		}
		long long _find(int l, int r, int i, int li, int ri) {
			if (lazy[i]) {
				if (li != ri) {
					lazy[2 * i + 1] += lazy[i];
					lazy[2 * i + 2] += lazy[i];
				}
				tree[i] += lazy[i];
				lazy[i] = 0;
			}
			if (l == li && r == ri) return tree[i];
			int mid = (li + ri) / 2;
			if (r <= mid) return _find(l, r, 2 * i + 1, li, mid);
			if (l >= mid + 1) return _find(l, r, 2 * i + 2, mid + 1, ri);
			return min(_find(l, mid, 2 * i + 1, li, mid), _find(mid + 1, r, 2 * i + 2, mid + 1, ri));
		}
		long long _upd(int l, int r, int d, int i, int li, int ri) {
			if (l == li && r == ri) {
				lazy[i] += d;
				return tree[i] + lazy[i];
			}
			int mid = (li + ri) / 2;
			long long left, right;
			if (r <= mid) {
				left = _upd(l, r, d, 2 * i + 1, li, mid);
				right = tree[2 * i + 2] + lazy[2 * i + 2];
			}
			else if (l > mid) {
				left = tree[2 * i + 1] + lazy[2 * i + 1];
				right = _upd(l, r, d, 2 * i + 2, mid + 1, ri);
			}
			else {
				left = _upd(l, mid, d, 2 * i + 1, li, mid);
				right = _upd(mid + 1, r, d, 2 * i + 2, mid + 1, ri);
			}
			tree[i] = min(left, right);
			return tree[i] + lazy[i];
		}
	public:
		LazyTree(vector<int>& arr) {
			n = arr.size();
			tree.resize(n * 4);
			lazy.resize(n * 4, 0);
			build(arr, 0, 0, n - 1);
		}
		long long find(int l, int r) { return _find(l, r, 0, 0, n - 1); }
		void upd(int l, int r, int d) { _upd(l, r, d, 0, 0, n - 1); }
};

아주 평범한 코드다.

lazy propagation에서 botoom-up 방식이 아예 불가능한 것은 아니다.

https://codeforces.com/blog/entry/18051

여기에 짧고 간결한 코드가 나와있다. 그런데 내 생각에 수정해서 쓰기는 긴 코드쪽이 조금 더 쉽지 않을까 생각한다.

예제는

https://www.acmicpc.net/problem/10999 (sum 세그먼트 트리)

https://www.acmicpc.net/problem/12844 (xor 세그먼트 트리)

https://codeforces.com/contest/52/problem/C (min 세그먼트 트리)

같은 것이 있다.