[BOJ] 13309 트리

문제 링크: https://www.acmicpc.net/problem/13309

13309번: 트리

문제 요약: 정점이 20만개인 트리가 주어진다. 이때 어떤 정점 a, b를 잇는 경로가 있는 지에 대한 쿼리 20만개에 응답해야 한다. 만약 쿼리가 원한다면, 응답 결과에 따라 a나 b의 부모와의 간선을 끊어야 할 수 있다. 

문제 풀이: 트리에서 어떤 두 정점 사이의 경로는 유일하다. 그러니까 다른 말로 하면 쿼리는 a와 b 사이의 경로 중에 있는 간선 중 끊어진 것이 하나라도 존재하는지 응답하는 쿼리로 환원될 수 있다. 이를 위해 다음과 같은 작업을 해 줄 수 있다.

1. 어떤 정점의 서브트리에 속한 정점들에 올라갈 수 있는 최대의 높이를 기록하기.
2. 두 정점 a,b의 경로에서 지나야 하는 가장 높은 정점의 높이를 찾아내기.

1번 작업과 2번 작업을 빠르게 할 수 있다면, 문제는 다음과 같이 풀 수 있다. 어떤 정점 x와 그 부모 정점을 잇는 간선이 제거되면, x의 서브트리에 있는 모든 정점들은 올라갈 수 있는 높이가 아무리 높아도 최대 x의 높이로 제한된다. 따라서 어떤 두 정점 a,b의 경로에서 지나야 하는 가장 높은 정점을 t라고 하자. 이때 경로가 존재하려면 a에서 t로 올라갈 수 있어야 하고, b에서 t로 올라갈 수 있어야 할 것이다. 따라서 a와 b에서 올라갈 수 있는 최대 높이가 둘 다 t의 높이보다 높으면, 경로가 존재함을 알 수 있다.

그러면 문제는 이 두 작업을 빠르게 하는 방법을 찾는 것이다. 다행히도, 1번 작업은 세그먼트 트리 레이지 프로파게이션을 통해서 빠르게 해 줄 수 있다. 세그먼트 트리에 날려야 하는 쿼리가 항상 점 쿼리여서 쉽게 코딩할 수 있다. 트리를 dfs ordering 순으로 펴 준 다음 서브트리 사이즈를 기록하면 어떤 서브트리 전체에 업데이트 쿼리를 날려줄 수 있다. 2번 작업은 lca(a,b)임이 당연하고, sparse table로 쉽게 구할 수 있다. 두 연산 모두 시간복잡도가 $O(lg n)$ 이므로, 전체 시간복잡도는 $O(q lg n)$ 이 된다. 

코드는 다음과 같다.

#define _CRT_SECURE_NO_WARNINGS
#include <vector>
#include <algorithm>
#include <cstdio>
#define MAXN 200005

using namespace std;

int n, q, pr;
int par[MAXN][20], depth[MAXN], sz[MAXN], idx[MAXN], used[MAXN];
vector<int> child[MAXN];

int check[MAXN];
int dfs(int x, int dep) {
	idx[x] = ++pr;
	int ssz = 1;
	depth[x] = dep;
	for (int there : child[x]) {
		if (check[there]) continue;
		ssz += dfs(there, dep + 1);
	}
	return sz[x] = ssz;
}

void pre_process() {
	dfs(1, 1);
	for (int i = 1; i <= 19; i++) {
		for (int j = 1; j <= n; j++) {
			par[j][i] = par[par[j][i - 1]][i - 1];
		}
	}
}

int lca(int a, int b) {
	if (depth[a] > depth[b]) swap(a, b);
	for (int i = 19; i >= 0; i--) {
		if (depth[b] - (1 << i) >= depth[a])
			b = par[b][i];
	}
	if (a == b) return a;
	for (int i = 19; i >= 0; i--) {
		if (par[a][i] != par[b][i]) {
			a = par[a][i];
			b = par[b][i];
		}
	}
	return par[a][0];
}

typedef struct SegmentTree {
	int n;
	vector<int> tree, lazy;
	SegmentTree(int _n) : n(_n), tree(_n * 4 + 10), lazy(_n * 4 + 10) {}
	void lazy_update(int nd, int st, int en) {
		if (lazy[nd] == -1) return;
		if (st == en) 
			tree[nd] = max(tree[nd], lazy[nd]);
		else {
			lazy[nd * 2] = max(lazy[nd * 2], lazy[nd]);
			lazy[nd * 2 + 1] = max(lazy[nd * 2 + 1], lazy[nd]);
		}
		lazy[nd] = -1;
	}
	void update(int nd, int st, int en, int le, int ri, int val) {
		lazy_update(nd, st, en);
		if (en < le || ri < st) return;
		if (le <= st && en <= ri) {
			lazy[nd] = max(lazy[nd], val);
			return;
		}	
		int mid = (st + en) >> 1;
		update(nd * 2, st, mid, le, ri, val);
		update(nd * 2 + 1, mid + 1, en, le, ri, val);
	}
	void update(int le, int ri, int val) { update(1, 0, n - 1, le, ri, val); }
	int query(int nd, int st, int en, int pos) {
		lazy_update(nd, st, en);
		if (st == en && st == pos) return tree[nd];
		else if (pos < st || en < pos) return -1;
		int mid = (st + en) >> 1;
		if (pos <= mid) return query(nd*2, st, mid, pos);
		else return query(nd * 2 + 1, mid + 1, en, pos);
	}
	int query(int pos) { return query(1, 0, n - 1, pos); }
}ST;

int main() {
	scanf("%d %d", &n, &q);
	for (int i = 2; i <= n; i++)
		scanf("%d", &par[i][0]);
	for (int i = 2; i <= n; i++)
		child[par[i][0]].push_back(i);
	pre_process();
	ST S(MAXN);
	while (q--) {
		int b, c, d, r;
		scanf("%d %d %d", &b, &c, &d);
		int k = lca(b, c);
		if (depth[k] >= max(S.query(idx[b]), S.query(idx[c]))) r = 1;
		else r = 0;
		printf("%s\n", r ? "YES" : "NO");
		if (d == 0) continue;
		if (r == 0) swap(b, c);
		if (used[b]) continue;
		used[b] = 1;
		S.update(idx[b], idx[b] + sz[b] - 1, depth[b]);
	}
	return 0;
}