[백준] 11437번 LCA - C/C++

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

11437번: LCA

문제

N(2 ≤ N ≤ 50,000) 개의 정점으로 이루어진 트리가 주어진다. 트리의 각 정점은 1번부터 N번까지 번호가 매겨져 있으며, 루트는 1번이다.

두 노드의 쌍 M(1 ≤ M ≤ 10,000)개가 주어졌을 때, 두 노드의 가장 가까운 공통 조상이 몇 번인지 출력한다.

입력

첫째 줄에 노드의 개수 N이 주어지고, 다음 N-1개 줄에는 트리 상에서 연결된 두 정점이 주어진다. 그다음 줄에는 가장 가까운 공통 조상을 알고 싶은 쌍의 개수 M이 주어지고, 다음 M개 줄에는 정점 쌍이 주어진다.

출력

M개의 줄에 차례대로 입력받은 두 정점의 가장 가까운 공통 조상을 출력한다.

---

문제에서 알 수 있듯이 LCA, 공통 조상을 찾는 문제이다.

node 들의 정보가 들어오면 인접 리스트를 만들어서 트리의 정보를 저장한다. 그리고 make_depth()를 통해 함수DFS 를 이용하여 트리를 순회하면서 각 노드들의 깊이 정보를 저장해둔다. 저장해둔 정보를 바탕으로 connection() 함수를 통해 각 노드의 2^k 번째 부모를 저장해둔다. 그다음 LCA() 함수를 통해 깊이가 더 깊은 노드를 다른 노드의 높이와 맞는 부모로 바꿔 두 노드를 같은 높이로 만들어 준다. 그다음 트리의 루트(보다 위)부터 내려오면서 공통 조상이 달라지는 순간 그 전 조상이 LCA 가 된다. 

#include <iostream>
#include <vector>
#include <cmath>
#include <cstring>

#define MAX 50001
#define MAX_H (int)floor(log2(MAX)) 
using namespace std;
int N, M;
vector<int> adj[MAX];
int parent[MAX][17];
int depth[MAX];

void make_depth(int cur){ // depth 정보를 만드는 함수
	for (int next : adj[cur]) { // range based for문
		if (depth[next] == -1) {
			parent[next][0] = cur;
			depth[next] = depth[cur] + 1;
			make_depth(next);
		}
	}
}

void connection() {
	for (int k = 1; k <= MAX_H; k++) {
		for (int cur = 1; cur <= N; cur++) {
			parent[cur][k] = parent[parent[cur][k - 1]][k - 1];
		}
	}
}

int LCA(int u, int v) {
	if (depth[u] < depth[v]) { 
		int tmp = u;
		u = v;
		v = tmp;
	}

	int diff = depth[u] - depth[v];
	for (int i = 0; diff != 0; i++) {
		if ((diff & 1) == 1)
			u = parent[u][i];
		diff = diff >> 1; 
	}

	if (u != v) {
		for (int i = MAX_H; i >= 0; i--) {
			if (parent[u][i] != 0 && parent[u][i] != parent[v][i]) {
				u = parent[u][i];
				v = parent[v][i];
			}
		}
		u = parent[u][0];
	}
	return u;
}

int main() {
	cin >> N;
	memset(depth, 0, sizeof(depth));
	memset(parent, -1, sizeof(parent));
	depth[1] = 0;
	for (int i = 0; i < N-1; i++) {
		int a, b;
		cin >> a >> b;
		adj[a].push_back(b);
		adj[b].push_back(a);
	}
	make_depth(1);
	connection();
	cin >> M;
	for (int i = 0; i < M; i++) {
		int u, v;
		cin >> u >> v;
		cout << LCA(u, v) << "\n";
	}
}

---

SWEA의 문제 중에 공통조상 찾기 문제가 2문제나 있어서 아침에 LCA 문제를 공부하기 위해서 백준에서 찾아 풀었다. 모든 조상을 기록하지 않고 2^k 번째 조상만 저장 해두고 2^(k+k) 조상은 2^k 번째 조상의 2^k번째 조상인 것이나 비트 마스킹을 이용하여 깊이를 맞추는 등 접해보지 못한 아이디어들이 많아 신기한 것 같다. 이 알고리즘이 잊혀 갈 때쯤 1761번: 정점들의 거리 / 15480번: LCA와 쿼리 문제를 풀어봐야겠다!