BOJ3584: 가장 가까운 공통 조상

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

  • 이제는 쭉 구현할 수 있는 LCA 문제.
  • Tree DP의 기본이 되는 문제인 만큼, 손에 익혀두자.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#include <bits/stdc++.h>
using namespace std;

class LCA {
  int N;
  vector<int> depth, indegree;
  vector<vector<int>> parent, adj;
public:
  LCA(int n) {
    N = n;
    depth.resize(n+1);
    indegree.resize(n+1);
    parent.resize(n+1, vector<int>(21));
    adj.resize(n+1);
  }

  void addEdge(int parent, int son);
  int findRoot();
  void dfs(int n, int p, int d);
  void fillTable();
  int lca(int a, int b);
};

void LCA::addEdge(int parent, int son) {
  indegree[son]++;
  adj[parent].push_back(son);
}

int LCA::findRoot() {
  for(int i = 1; i <= N; i++)
    if(indegree[i] == 0)
      return i;
}

void LCA::dfs(int n, int p, int d) {
  parent[n][0] = p;
  depth[n] = d;
  for(int m: adj[n])
    dfs(m, n, d+1);
}

void LCA::fillTable() {
  for(int j = 1; j < 21; j++)
    for(int i = 1; i <= N; i++)
      parent[i][j] = parent[parent[i][j-1]][j-1];
}

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

int main() {
  ios::sync_with_stdio(false);
  int TC;
  cin >> TC;
  while(TC--) {
    int n, a, b;
    cin >> n;
    LCA lca(n);
    for(int i = 0; i < n-1; i++) {
      cin >> a >> b;
      lca.addEdge(a, b);
    }
    int root = lca.findRoot();
    lca.dfs(root, 0, 0);
    lca.fillTable();
    cin >> a >> b;
    printf("%d\n", lca.lca(a,b));
  }
}