세그먼트 트리

학교 동아리에서 매주 배우는 알고리즘을 정리해 볼 것이다. 그 뒤 동아리에서 추천해준 백준문제를 풀어 볼 것 이다.

사용하는 경우

1. 수열에서 임의의 구간에 대한 최대값, 최소값, 합, 곱 등을 빠르게 구할 때 (세그먼트 트리)
2. 수열에서 임의의 구간에 대한 업데이트를 빠르게 할 때 (레이지 세그먼트 트리)

세그먼트 트리 모습

수열 [5,7,3,2,9,8] 을 합의 세그먼트 트리 모습으로 나타내면

이런 모습으로 보여진다. 리프 노드가 수열이고 다른 노드가 합을 나타탠다. 트리는 배열로 나타낼 수 있다. 인덱스 1 의 왼쪽 자식이 2 오른쪽 자식이 3 이런 식이다. 즉, node x의 왼쪽 자식 노드는 2x, 오른쪽 자식 노드는 2x+1이다.

위 트리를 배열로 나타내면 이렇게 나타낼 수 있다.

세그먼트 트리 사용

위 트리에서 인덱스 2 ~ 인덱스 5 까지의 합을 구할려면 어떻게 동작하는지 알아보자.

합 구하기

인덱스 2부터 5까지의 합은 7+3+2+9 이다. 노드 34에서 시작해서 한단계식 내려온다. 15로와서 15가 인덱스 1,2,3을 포함하니 한단계 더 내려온다. 인덱스 3은 멈추고 12가 인덱스 1,2 를 포함하므로 한단계 또 내려오게 된다. 19는 인덱스 4,5,6을 포함하고 있다. 11로가니 11이 인덱스 4,5를 모두 포함하고 있으므로 멈추고 이 값을 반환한다. 즉 7+3+11을 하게 된다. 이렇게 합을 구할 수가 있다.

최대, 최소 구하기

최대 최소를 구하는 것은 합의 결과를 더했던 칸에 합 말고 최대와 최소를 넣으면 된다.

업데이트

세그먼트 트리는 원하는 인덱스에 얼마의 값을 추가 할 수 있다. 예를 들어 인덱스 1~ 인덱스 4에 10을 추가해라 라고 할 수 있다. 15로 내려와 15가 인덱스 1~3을 포함하고 있으므로 15에 103을 더한다. 그리고 12와 3으로 내려와 3에 10을 추가하고 12에는 102를 추가한다. 마지막으로 5와 7에 각각 10을 추가하게 된다. 인덱스 4는 리프노드 까지 내려와 10을 추가하고 11과 19에도 각각 10을 추가하게 된다. 결국 34에 40을 추가하게 된다.

구현

내가 구현해볼 세그먼트 트리는 합을 보여줄 세그먼트 트리이다.

const int MAX = 1000000;
long long tree[MAX * 4 + 100]; // tree 구현할 배열
long long lazy[MAX * 4 + 100]; // 더해야할 수가 있는지 저장할 배열
long long a[MAX + 10];  // 입력받은 수열

1. 초기화

제일 처음 배열들로 세그먼트 트리를 만드는 코드부터 구현해야 한다.

void init(int node, int s, int e) {
	if (s == e) { //처음과 끝이 같을 때는 leaf 노드 
		tree[node] = a[s];
		return;
	}
	int mid = (s + e) / 2;
	init(node * 2, s, mid); //왼쪽 자식 초기화
	init(node * 2 + 1, mid + 1, e); //오른쪽 자식 초기화
	tree[node] = tree[node * 2] + tree[node * 2 + 1]; //leaf 노드까지 
  간다음 그 위는 합으로 초기화
}

2. Propagation

합을 구하는 문제나 업데이트 문제를 할 때 각 노드에 더할게 있는지 검사하는 함수이다.

void propagation(int node, int s, int e) { //더할꺼 있는 지 검사
	if (lazy[node] == 0) return; //지금 검사할 노드에 lazy에 아무것도 없으면 리턴
	tree[node] += lazy[node] * (e - s + 1); //자신의 총 자식수만큼 더함 
	if (s != e) { //s 랑 e 가다르면 더 밑에 내려가야하므로 lazy도 내림
		lazy[node * 2] += lazy[node];
		lazy[node * 2 + 1] += lazy[node];
	}
	lazy[node] = 0;
}

3. 합을 구하는 함수

어디서부터 어디서 까지 합을 구하라 라고 알려주면 수행할 함수.

long long query(int node, int s, int e, int l, int r) { 
	propagation(node, s, e); //더할게 있는지 검사
	if (r < s || e < l) return 0; // r이 start보다 작거나 l이 end보다 작은것은 예외처리
	if (l <= s && e <= r) return tree[node]; //모든 수의 합임
	int mid = (s + e) / 2;
	long long left = query(node * 2, s, mid, l, r);
	long long right = query(node * 2 + 1, mid + 1, e, l, r);
	return left + right;
}

4. 구간에 얼마만큼 더해야 할 함수

  구간 어디서부터 어디까지중 각각 몇을 더해라를 해결할 때 쓰는 함수

void update(int node, int s, int e, int l, int r, long long x) { 
	propagation(node, s, e);
	if (r < s || e < l)return; 
	if (l <= s && e <= r) { // l이 start보다 작고 r이 end보다 클때는 모두 다 더해야한다.
		lazy[node] += x;
		propagation(node, s, e);
		return;
	}
	int mid = (s + e) / 2;
	update(node*2, s, mid, l, r, x);
	update(node * 2 + 1, mid + 1, e, l, r, x);
	tree[node] = tree[node * 2] + tree[node * 2 + 1]; // 더한수에 대해 다시 합 구하기
}

정리

세그먼트 트리에 대해 공부해 보았다. 이 트리는 구간 합이나 빠르게 최소값 최대값을 구할때나 특정한 범위에 수를 더할 때 쓸 수 있는 것 같다. 더 많은 활용은 백준 문제를 풀어가면서 알아 가야겠다.