[백준] 14268 회사 문화 2 - C++

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

---

얼핏 보면 간단히 합세그만 구현해도 되는게 아닌가? 싶지만

문제에서 주는 조직도는 결국 트리 형태이고, 쿼리 또한 서브트리를 대상으로 한다.

하지만 세그트리는 선형 구조(배열)에서 사용할 수 있기 때문에 무지성 합세그구현은 못한다..

이럴땐 트리를 선형으로 펴주는 ETT를 사용하면 된다!

---

ETT(Euler Tour Technique) 란?

ETT란, 임의의 노드의 서브트리를 선형구조의 구간으로 표현할수 있게 하는 알고리즘(테크닉)이다.

세그트리와 함께 사용하면 서브트리의 합을 구할 수 있는 것이다!

출처 - kokodak tistory

 

(1번 노드가 루트 노드라고 생각하자.)

ETT는 기본적으로 DFS가 노드를 방문하는 순서에 따라 번호를 매긴다.

DFS를 수행하는 과정에서 직접 노드를 방문하는 가상의 커서가 있다고 생각해보자.

이때 가상의 커서가 i번째 노드에 최초로 진입한 시점을 in[i]이라고 하자.

또한 i번째 노드의 모든 자식의 탐색이 완료된 시점을 out[i]라고 하자.

위 그림에서 노드 당 in은 검은 글씨, out은 파란 글씨로 표현되어있다.

 

이렇게 방문 순서를 in과 out으로 표현하면 뭐가 좋느냐??

임의의 노드 i가 있을때, i를 루트로 삼는 서브트리의 모든 노드 j는

in[i] < in[j], out[j] <= out[i]를 만족한다.

한마디로 in[i] ~ out[i] 구간에는 i를 루트로 삼는 서브트리의 모든 노드가 포함된다!

 

이렇게 비선형인 트리구조를 in과 out을 통해 선형구조처럼 만들었으니 이제 세그트리를 사용할 수 있다!

---

부분 코드 설명

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

vector<ll> in(100'005);
vector<ll> out(100'005);

ll seg[100'000 * 4 + 16];
ll lazy[100'000 * 4 + 16];

vector<int> g[100'005];

선언부이다.. 문제조건에서 n이 100'000이라 했으니 약간 널널하게 배열 사이즈를 잡았다.

ETT를 위한 in과 out도 선언하고

서브트리를 대상으로 구간 업데이트도 해야하기 때문에 lazy prop용 배열도 만들어준다...

g(graph)는 조직도를 나타내는 원본 트리이다.

 

int cnt = 0;
void dfs(int cur) {
	in[cur] = ++cnt; //<1>
	for (auto& next : g[cur]) { //<2>
		dfs(next);
	}
	out[cur] = cnt; //<3>
}

ETT를 진행하는 DFS 함수이다. 전역변수 cnt를 사용해 방문순서를 기록한다.

<1> : cnt를 1 증가시키고, 그 값을 in에 저장한다.

<2> : 그냥 조직도 훑으면서 재귀 DFS 돌린다.

<3> : 해당 노드의 서브트리의 탐색이 종료된 시점에 실행된다. out을 cnt값으로 저장한다.

탐색이 완료된 시점을 저장하는것이기에 cnt를 증가시키지는 않는다.

 

세그트리는 그냥 lazy 합세그로 구현하면 된다.

void push(ll idx, ll left, ll right) {
	if (lazy[idx] != 0) {
		seg[idx] += lazy[idx] * (right - left + 1);
		if (left != right) {
			lazy[idx * 2] += lazy[idx];
			lazy[idx * 2 + 1] += lazy[idx];
		}
		lazy[idx] = 0;
	}
}
void update(ll idx, ll left, ll right, ll cleft, ll cright, ll cval) {
	push(idx, left, right);
	if (cright < left || right < cleft) {
		return;
	}
	if (cleft <= left && right <= cright) {
		lazy[idx] += cval;
		push(idx, left, right);
		return;
	}
	ll mid = (left + right) / 2;
	update(idx * 2, left, mid, cleft, cright, cval);
	update(idx * 2 + 1, mid + 1, right, cleft, cright, cval);
	seg[idx] = seg[idx * 2] + seg[idx * 2 + 1];
	return;
}
ll query(ll idx, ll left, ll right, ll qidx) {
	push(idx, left, right);
	if (qidx < left || right < qidx) {
		return 0;
	}
	if (left == right) {
		return seg[idx];
	}
	ll mid = (left + right) / 2;
	return query(idx * 2, left, mid, qidx) + query(idx * 2 + 1, mid + 1, right, qidx);
}

 

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	int n, m;
	cin >> n >> m;
	for (int i = 1; i <= n; i++) { //<1>
		int x;
		cin >> x;
		if (x != -1) g[x].push_back(i);
	}

	dfs(1); //<2>

	for (int i = 0; i < m; i++) {
		int command, a, b;
		cin >> command;
		if (command == 1) {
			cin >> a >> b;
			update(1, 1, n, in[a], out[a], b); //<3>
		}
		else {
			cin >> a;
			cout << query(1, 1, n, in[a]) << "\n"; //<4>
		}
	}
	return 0;
}

main이다. 

<1> : 조직도를 입력받는다. i번째 노드의 부모 노드 x가 주어지고, 1번노드는 -1이 주어진다.

이 점에 유의하며 코드를 적자.

<2> : ETT용 DFS를 돌린다. 루트노드의 번호인 1을 넣어주면 된다.

<3> : a번째 직원이 b만큼 칭찬을 받으면 a번째 노드를 루트로 하는 서브트리의 모든 노드가

b만큼 칭찬을 받게된다. 이 말은 즉슨 아까 말했듯 in[a] ~ out[a] 구간에 b만큼 더해주면 된다.

<4> : a번째 노드의 값을 출력하는 쿼리이다.

트리에서의 a번째 노드가 선형구조에서는 in[a]를 인덱스로 가지므로

in[a]의 값을 쿼리로 날리면 된다.

---

전체 코드

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

vector<ll> in(100'005);
vector<ll> out(100'005);

ll seg[100'000 * 4 + 16];
ll lazy[100'000 * 4 + 16];

vector<int> g[100'005];

int cnt = 0;
void dfs(int cur) {
	in[cur] = ++cnt;
	for (auto& next : g[cur]) {
		dfs(next);
	}
	out[cur] = cnt;
}
void push(ll idx, ll left, ll right) {
	if (lazy[idx] != 0) {
		seg[idx] += lazy[idx] * (right - left + 1);
		if (left != right) {
			lazy[idx * 2] += lazy[idx];
			lazy[idx * 2 + 1] += lazy[idx];
		}
		lazy[idx] = 0;
	}
}
void update(ll idx, ll left, ll right, ll cleft, ll cright, ll cval) {
	push(idx, left, right);
	if (cright < left || right < cleft) {
		return;
	}
	if (cleft <= left && right <= cright) {
		lazy[idx] += cval;
		push(idx, left, right);
		return;
	}
	ll mid = (left + right) / 2;
	update(idx * 2, left, mid, cleft, cright, cval);
	update(idx * 2 + 1, mid + 1, right, cleft, cright, cval);
	seg[idx] = seg[idx * 2] + seg[idx * 2 + 1];
	return;
}
ll query(ll idx, ll left, ll right, ll qidx) {
	push(idx, left, right);
	if (qidx < left || right < qidx) {
		return 0;
	}
	if (left == right) {
		return seg[idx];
	}
	ll mid = (left + right) / 2;
	return query(idx * 2, left, mid, qidx) + query(idx * 2 + 1, mid + 1, right, qidx);
}
int main() {

	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	int n, m;
	cin >> n >> m;
	for (int i = 1; i <= n; i++) {
		int x;
		cin >> x;
		if (x != -1) g[x].push_back(i);
	}

	dfs(1);

	for (int i = 0; i < m; i++) {
		int command, a, b;
		cin >> command;
		if (command == 1) {
			cin >> a >> b;
			update(1, 1, n, in[a], out[a], b);
		}
		else {
			cin >> a;
			cout << query(1, 1, n, in[a]) << "\n";
		}
	}
	return 0;
}

---

ETT.. 정말 간단하지만 재밌는 아이디어 같다..

---