백준 11437번 LCA C++ LCA, 최소 공통 조상

리뷰

새로운 알고리즘인 LCA를 배우게 되었다, 로직이 매우 복잡했지만 세그먼트 트리처럼 외워서 쓴다면 괜찮을 것 같다.

 

문제 풀이

1. nodes를 최대 노드의 개수로, LOG를 최대 노드의 개수를 log2를 취한 수보다 크게 설정해 준다.
2. tree벡터 배열에 인접 리스트를 저장하고, depth에 각 노드의 깊이, parent에 각 노드의 부모 정보를 저장한다.
3. dfs 함수를 통해 인접리스트를 사용하여 자신의 바로 윗 부모와 depth를 알 수 있다.
4. process 함수를 통해 모든 노드의 조상에 대해 탐색이 가능하다, 조상이 없을 경우 -1
5. m개의 쿼리문을 받아와 u와 v의 최소 공통 부모를 lca 함수를 통해 구해준 후 출력해 준다.

 

참고 사항

자세한 내용은 정답 코드의 주석 참고

 

 

정답 코드

#include<iostream>
#include<vector>
#include<algorithm>

using namespace std;

const int nodes = 50001; // 최대 노드의 개수
const int LOG = 20; // log2 최대 노드의 개수 보다 큰 값으로 설정

vector<int> tree[nodes]; // 인접 리스트 생성용
int depth[nodes]; // 노드의 깊이를 나타낼 배열
int parent[nodes][LOG]; // 노드의 조상을 나타낼 배열

int n, m;

void dfs(int node, int par, int dep) { // 인접리스트를 활용하여 첫번째 조상 구하기
	depth[node] = dep; // 노드의 깊이 초기화
	parent[node][0] = par; // 노드의 첫번째 부모 초기화

	for (int child : tree[node]) { // 인접 리스트 순회
		if (child != par) { // 무한루프 방지
			dfs(child, node, dep + 1); // 재귀 실행
		}
	}
}

void process() { // 모든 조상 탐색
	for (int j = 1; j < LOG; j++) { 
		for (int i = 1; i <= n; i++) {
			if (parent[i][j - 1] != -1) { // 더 이상 조상이 없는 경우가 아니라면
				parent[i][j] = parent[parent[i][j - 1]][j - 1]; // 그 다음 조상은 내 부모의 조상
			}
		}
	}
}

int lca(int u, int v) { // u와 v의 조상을 이진 탐색
	if (depth[u] < depth[v]) swap(u, v); // u가 더 크게 유지
	int diff = depth[u] - depth[v]; // u와 v의 깊이 차이를 구한다.

	for (int i = 0; i < LOG; i++) { // 비트마스킹을 통한 깊이 맞추기
		if ((diff >> i) & 1) u = parent[u][i]; // i번째 비트가 1일 경우 2^i 만큼 깊이를 올린다.
	}

	if (u == v) return u; // v가 u의 조상일 경우 리턴

	for (int i = LOG - 1; i >= 0; i--) { // 동시에 깊이를 조정하며 일치하는 부모를 찾는다.
		if (parent[u][i] != parent[v][i]) { // 부모가 다르면 깊이 이동
			u = parent[u][i];
			v = parent[v][i];
		}
	}
	return parent[u][0]; // 바로 위의 부모가 최소 공통 부모가 된다.
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	cin >> n;
	for (int i = 0; i < n - 1; i++) { // 인접리스트 초기화
		int a, b;  cin >> a >> b;
		tree[a].push_back(b);
		tree[b].push_back(a);
	}
	dfs(1, -1, 0); // 루트부터 dfs 탐색
	process();
	cin >> m;
	while (m--) {
		int a, b; cin >> a >> b;
		cout << lca(a, b) << "\n";
	}
}