루트가 1번인 트리에서 두 노드의 가장 가까운 공통 조상을 구하는 문제였다.
노드 수가 최대 50,000개고, 쿼리도 10,000개나 돼서 효율적인 풀이가 필요했다.
처음엔 그냥 각 노드의 부모를 따라 올라가면서 공통 조상을 찾는 식으로 풀었다.
그랬더니 당연하게도 시간 초과가 났다. 노드 수나 쿼리 수 생각하면 이건 O(N)짜리 루프를 M번 도는 꼴이라 당연한 결과였다.
그래서 이 문제는 LCA 전용 알고리즘을 써야겠다고 판단했다.
Binary Lifting이란 걸 처음 써봤는데, 생각보다 구현이 복잡하진 않았다.
처음엔 희소 배열의 의미를 잘 몰라서 2차원 배열을 어떻게 채워야 할지 좀 헷갈렸지만, 한 번 구현해보니까 금방 익숙해졌다.
풀이 방식
- DFS로 각 노드의 깊이와 1단계 부모를 구했다.
- parent[node][k]를 node의 2^k번째 조상으로 설정해서 희소 배열을 채웠다.
- LCA 쿼리 시엔 먼저 두 노드의 깊이를 맞췄고,
- 동시에 위로 올라가면서 처음으로 만나는 공통 조상을 찾았다.
쿼리 하나 처리하는 데 O(log N)이고, 전체적으로 M log N에 문제를 풀 수 있었다.
#include
#include
#include
#define MAXN 50001
#define LOGN 17 // 2^16 = 65536 > 50000
int N, M;
int depth[MAXN];
int parent[MAXN][LOGN];
int adj[MAXN][MAXN / 10];
int adj_size[MAXN];
int visited[MAXN];
void dfs(int node, int p, int d) {
visited[node] = 1;
depth[node] = d;
parent[node][0] = p;
for (int i = 0; i < adj_size[node]; i++) {
int next = adj[node][i];
if (!visited[next]) {
dfs(next, node, d + 1);
}
}
}
void preprocess() {
for (int j = 1; j < LOGN; j++) {
for (int i = 1; i <= N; i++) {
if (parent[i][j - 1] != 0)
parent[i][j] = parent[parent[i][j - 1]][j - 1];
}
}
}
int lca(int a, int b) {
if (depth[a] < depth[b]) {
int temp = a;
a = b;
b = temp;
}
int diff = depth[a] - depth[b];
for (int i = 0; i < LOGN; i++) {
if ((diff >> i) & 1)
a = parent[a][i];
}
if (a == b) return a;
for (int i = LOGN - 1; i >= 0; i--) {
if (parent[a][i] != 0 && parent[a][i] != parent[b][i]) {
a = parent[a][i];
b = parent[b][i];
}
}
return parent[a][0];
}
int main() {
scanf("%d", &N);
for (int i = 0; i < N - 1; i++) {
int u, v;
scanf("%d %d", &u, &v);
adj[u][adj_size[u]++] = v;
adj[v][adj_size[v]++] = u;
}
dfs(1, 0, 0);
preprocess();
scanf("%d", &M);
for (int i = 0; i < M; i++) {
int a, b;
scanf("%d %d", &a, &b);
printf("%d\n", lca(a, b));
}
return 0;
}
정리
- 트리에서 조상 찾는 문제는 무조건 LCA를 떠올려야 한다.
- Binary Lifting을 쓰면 O(log N)으로 조상 탐색 가능하다.
- 희소 배열 전처리만 잘 해두면, 쿼리 처리도 깔끔하게 된다.
- DFS로 깊이랑 부모를 구하는 게 핵심이다.
처음엔 무작정 올라가는 방식으로 접근했다가 시간 초과가 났다.
시간 복잡도 계산을 안 하고 무식하게 풀었던 게 문제였다.
문제를 보고 LCA를 떠올리지 못한 것도 내 실수였다...ㅎ
이번 기회에 Binary Lifting이 어떤 방식으로 작동하는지 정확히 이해할 수 있었다.
다음에 트리 + 조상 조합의 문제가 나오면 바로 이 방식부터 떠올리면 될 것 같다.
트리에서 조상 찾기 -> LCA 먼저 떠올리기!! 하나 알아갑니당~