[P2] 백준 15480번 LCA와 쿼리 C++ LCA, 최소 공통 조상, 트리

리뷰

 

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

LCA의 심화 문제, 루트가 매번 바뀔때마다의 처리하는 아이디어에 배우게 된 문제였다.

 

전역 변수

• MAX_N : 정점의 최대 개수를 저장할 정수형 상수 변수
• LOG : DP에서 2의 N승으로 빠르게 탐색하기 위해 사용할 정수형 상수 변수
• n, m : 문제에서 주어진 정점의 개수 n, 쿼리의 개수 m
• depth : 각 노드의 초기 깊이를 저장할 정수형 배열, 크기는 MAX_N으로 초기화
• parents : 각 노드의 부모 정보를 저장해 둔 정수형 배열, 크기는 MAX_N * LOG로 초기화
• lst : 각 정점간의 간선을 인접리스트로 저장할 정수형 벡터, 크기는 MAX_N으로 초기화

 

함수

1. dfs

void dfs(int node, int par, int dep)

 

트리의 관계를 설정하고 depth와 parents배열의 초깃값을 저장할 함수

1. 매개변수로 현재 노드 번호 node, 부모 노드 번호 par, 현재 노드의 깊이 dep을 받아온다.
2. 현재 노드의 depth를 dep으로 저장하고, parents의 바로 위 조상을 par로 저장한다.
3. 인접리스트를 순회하며 각 자식에게 재귀를 진행해 준다.
4. 재귀를 진행할땐 node가 child로, par가 node로, dep를 1 증가시켜 주어야 한다.

 

2. preprocess

void preprocess()

 

parents 정보를 마저 채워주기 위한 전처리 함수

1. j를 1 ~ LOG - 1까지, i를 1 ~ n까지 2중 for문을 개행해 준다.
2. 현재 노드의 2^j - 1번째 조상이 존재할 경우 j번째 조상은 내 2^j - 1번째 조상의 2^j - 1번째 조상으로 저장한다.

3. lca

int lca(int u, int v)

 

두 노드 u, v의 최소 공통 조상을 구하는 함수

1. 우선 탐색을 편하게 하기 위해 u와 v의 depth중 v가 더 깊다면 u와 v를 바꿔준다.
2. u와 v의 depth차이를 정수형 변수 diff에 저장해 준다.
3. diff의 비트를 순회하며 만약 i번째 비트가 1이라면 u를 u의2^i번째 조상으로 바꾸어 준다.
4. 해당 작업을 마치면 u와 v의 깊이가 동일해 졌다, 만약 둘이 같은 노드라면 그대로 u를 리턴해 준다.
5. 같은 노드가 아니라면 LCA의 바로 밑까지 두 노드를 올려주는 작업을 수행한다.
6. LOG - 1 ~ 0까지 탐색을 진행하고 만약 2^i번째 부모가 서로 다르다면 둘다 2^i만큼 올려준다.
7. 해당 작업을 마친 후 u의 바로 위 조상을 리턴하면 LCA를 구할 수 있다.

4. newRoot

int newRoot(int r, int u, int v)

 

루트가 r일때 u와 v의 공통 조상을 구하는 함수

1. r, u // u, v // r, v각각 lca의 인자로 넣어 ru, uv, rv 노드를 구해준다.
2. 세 노드를 모두 비교하여 깊이가 가장 깊은 노드를 리턴해 준다.

 

문제풀이

1. n값을 입력 받고, n - 1개의 간선 정보를 입력받아 양방향 인접 리스트를 만들어 준다.
2. dfs함수를 통해 depth 배열을 초기화 하고, parents 배열의 밑작업을 진행해 준다.
3. preprocess함수를 통해 parents 배열의 정보를 완성시켜 준다.
4. m을 입력 받고, m번의 while문을 개행해 준다.
5. r, u, v를 입력 받고, newRoot 함수에 매개변수로 전달하여 리턴된 값을 출력해 준다.

 

참고 사항

depth와 parent 정보를 DFS를 통해 계산할 때는 기본적으로 루트가 1인 경우로 계산된다.

하지만 이 정보는 트리의 기본적인 구조에 대한 정보이기 때문에, 루트가 변경되더라도 LCA를 계산하는 데 문제가 없다.

트리의 구조(간선 정보)는 루트가 무엇이든 변하지 않는다.

즉, 두 노드 사이의 부모-자식 관계나 그들이 어느 높이에 있는지(깊이)는 트리의 루트가 무엇인지와 무관하다.

따라서 DFS로 depth와 parent를 계산할 때, 루트가 1이라고 가정하고 계산한 것이 트리 전체의 구조를 반영한 것이기 때문에, 루트가 다른 노드로 변경되더라도 이 정보는 유효하다.

 

정답 코드

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

using namespace std;

const int MAX_N = 100001;
const int LOG = 18;

int n, m;
int depth[MAX_N];
int parents[MAX_N][LOG];

vector<int> lst[MAX_N];

void dfs(int node, int par, int dep) {
	depth[node] = dep;
	parents[node][0] = par;

	for (int child : lst[node]) {
		if (child != par) dfs(child, node, dep + 1);
	}
}

void preprocess() {
	for (int j = 1; j < LOG; j++) {
		for (int i = 1; i <= n; i++) {
			if (parents[i][j - 1] != -1) parents[i][j] = parents[parents[i][j - 1]][j - 1];
		}
	}
}

int lca(int u, int v) {
	if (depth[u] < depth[v]) swap(u, v);

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

	if (u == v) return u;

	for (int i = LOG - 1; i >= 0; i--) {
		if (parents[u][i] != parents[v][i]) {
			u = parents[u][i];
			v = parents[v][i];
		}
	}

	return parents[u][0];
}

int newRoot(int r, int u, int v) {
	int ru = lca(r, u);
	int uv = lca(u, v);
	int rv = lca(r, v);

	if (depth[ru] > depth[uv] && depth[ru] > depth[rv]) return ru;
	if (depth[uv] > depth[ru] && depth[uv] > depth[rv]) return uv;
	return rv;
}

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;
		lst[a].push_back(b);
		lst[b].push_back(a);
	}

	dfs(1, -1, 0);
	preprocess();

	cin >> m;
	while (m--) {
		int r, u, v; cin >> r >> u >> v;
		cout << newRoot(r, u, v) << "\n";
	}
}