[BOJ 15480] LCA와 쿼리(c++)

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

문제 개요

트리에서 LCA를 구하는 문제인데, 각 쿼리마다 root가 다릅니다. N <= 100,000이니 이진 점프로 빠르게 LCA를 구해주어야 합니다.

 (1 <= N, M <= 100,000)

 (1 <= r, u, v <= N)

접근

처음에는 오프라인 쿼리로 같은 루트에서 LCA를 구하는 쿼리끼리 모아서 루트 별로 새로 level과 parent를 구해서 풀었는데, 시간 초과가 나더군요. 

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int n, m;
struct node {
	int r, x, y, idx;
};
int res[100001];
vector<node> queries;
vector<int> v[100001];

int level[100001];
int par[100001][18];

bool cmp(node a, node b) {
	return a.r < b.r;
}

void dfs(int x, int pre) {
	for (int y : v[x]) {
		if (y == pre) continue;
		level[y] = level[x] + 1;
		par[y][0] = x;
		dfs(y, x);
	}
}

void init(int root) {
	level[root] = 1;
	dfs(root, 0);
	for (int k = 1; k < 18; k++) {
		for (int i = 1; i <= n; i++) {
			par[i][k] = par[par[i][k - 1]][k - 1];
		}
	}
}

int lca(int x, int y) {
	if (level[x] > level[y]) swap(x, y);
	for (int i = 17; i >= 0; i--) {
		if (level[y] - level[x] >= (1 << i)) {
			y = par[y][i];
		}
	}
	if (x == y) return x;
	for (int i = 17; i >= 0; i--) {
		if (par[x][i] != par[y][i]) {
			x = par[x][i];
			y = par[y][i];
		}
	}
	return par[x][0];
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	cin >> n;
	for (int i = 1; i < n; i++) {
		int x, y;
		cin >> x >> y;
		v[x].push_back(y);
		v[y].push_back(x);
	}
	cin >> m;
	queries.reserve(m);
	for (int i = 0; i < m; i++) {
		int r, x, y;;
		cin >> r >> x >> y;
		queries.push_back({ r,x,y,i });
	}
	sort(queries.begin(), queries.end(), cmp);
	int root = 0;
	for (node q : queries) {
		if (root != q.r) {
			root = q.r;
			init(root);
		}
		res[q.idx] = lca(q.x, q.y);
	}
	for (int i = 0; i < m; i++) {
		cout << res[i] << '\n';
	}
	return 0;
}

 init()함수의 시간 복잡도는 O(n log n)입니다. 모든 쿼리가 루트가 다 다르면 최악의 경우 init()함수로만 O(m * n log n)이니 n, m이 10만까지 들어올 수 있어서 시간초과가 납니다. root를 1로 두고 init()을 한 번만 해야겠다라 생각하고, 아이디어를 생각해보았습니다. root를 r로 두는 트리에서 u, v의 LCA를 구하면 다음과 같은 가능성이 있습니다.

 

 1. u - v의 경로에 r이 포함되는 경우

이 경우는 lca(u, v)는 r이 되야합니다.

 2. u - v 경로에 r이 포함되지 않는 경우

이 경우는 u - v가 지나는 경로들 중에 r과 가장 가까운 노드가 LCA가 되겠습니다. 

문제 풀이

먼저 u - v의 경로에 r이 포함되는 경우는 u와 v 중 반드시 하나는 r의 자손 노드일 것 입니다. 이 가정에서 lca(u, v)의 값이 반드시 r이 나오기 위해서는 lca(u, r)과 lca(v, r) 중 level이 더 높은 값을 고르면 됩니다. u, v 중 r의 자손 노드가 아닌 노드와의 LCA는 반드시 r보다 level이 높기 때문입니다. 이제 경로에 r이 포함되지 않는 경우를 살펴 보겠습니다. u, v가 모두 r의 자손 노드인 경우가 있는데, 이 경우는 lca(u, v)가 경로 중에서 r과 가장 가까운 노드일 것 입니다. 그리고 반대의 경우엔, lca(u, v)와 lca(r, u), lca(r, v) 중 level이 가장 높은 노드가 가까운 노드입니다.

따라서, lca(u, v), lca(u, r), lca(v, r) 중 level이 가장 높은 노드가 항상 root가 r인 트리에서의 u - v의 lca입니다.

코드

#include <iostream>
#include <vector>
using namespace std;

int n, m;
vector<int> v[100001];
int level[100001];
int parent[100001][18];

void dfs(int x, int pre) {
	for (int y : v[x]) {
		if (y == pre) continue;
		level[y] = level[x] + 1;
		parent[y][0] = x;
		dfs(y, x);
	}
}

void init() {
	level[1] = 1;
	dfs(1, 0);
	for (int k = 1; k < 18; k++) {
		for (int i = 1; i <= n; i++) {
			parent[i][k] = parent[parent[i][k - 1]][k - 1];
		}
	}
}

int lca(int x, int y) {
	if (level[x] > level[y]) swap(x, y);
	for (int i = 17; i >= 0; i--) {
		if (level[y] - level[x] >= (1 << i)) {
			y = parent[y][i];
		}
	}
	if (x == y) return x;
	for (int i = 17; i >= 0; i--) {
		if (parent[x][i] != parent[y][i]) {
			x = parent[x][i];
			y = parent[y][i];
		}
	}
	return parent[x][0];
}

int check(int a, int b, int c) {
	int ret = level[a] > level[b] ? a : b;
	ret = level[ret] > level[c] ? ret : c;
	return ret;
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	cin >> n;
	for (int i = 1; i < n; i++) {
		int x, y;
		cin >> x >> y;
		v[x].push_back(y);
		v[y].push_back(x);
	}
	init();
	cin >> m;
	for (int i = 0; i < m; i++) {
		int r, x, y;
		cin >> r >> x >> y;
		int a = lca(r, x);
		int b = lca(r, y);
		int c = lca(x, y);
		int res = check(a, b, c);
		cout << res << '\n';
	}

	return 0;
}