백준 13511번 트리와 쿼리 2 C++ 최소 공통 조상, LCA

리뷰

 

치열한 전투의 흔적이 보이십니까?

LCA의 기본기가 있다면 쉽게 풀 수 있을 거라 생각했는데, 문제가 친절하지 않아 고생한 문제

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

 

문제 풀이

1. input 함수에서 n값을 받고, 각 노드간의 간선 n - 1개를 입력 받는다.
2. solution함수에서 dfs 함수와 pre_process 함수를 돌려주어 각 노드의 부모, 깊이, 비용값을 구해주었다.
3. dfs 함수에서는 각 함수의 깊이 및 루트 노드로부터의 비용, 자신의 바로 위 부모를 구해 각 배열에 저장해줬다.
4. pre_process 함수에서는 각 노드의 2^i 번째 상위 부모를 찾아 부모 배열을 마저 초기화 해주었다.
5. query 함수에서 m값을 받고, 입력된 m개의 쿼리를 처리해 주었다.
6. 쿼리의 인덱스가 1인 경우, lca를 구해준 뒤 각 노드에서 lca까지의 비용을 각각 구한 뒤 더해주었다.
7. 인덱스가 2인 경우, k의 위치가 왼쪽에서 멈추는지, 오른쪽에서 멈추는지를 찾아 k위치의 노드를 출력해 주었다.

 

참고 사항

각 간선의 비용이 최대 100만, 간선의 개수는 최대 10만이므로 cost배열은 long long 타입으로 설정해 주어야 한다.

입력값이 많기 때문에 입력받는 시간을 최대한 줄여주어야 한다.

ios::sync_with_stdio(0);
cin.tie(0);

 

main 함수에 위 내용을 꼭 기재해 주도록 하자

 

정답 코드

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

using namespace std;

const int MAX_N = 100001; // 1 based index이므로 최대 노드의 개수
const int LOG = 20; // Log2(MAX_N)보다 크게 설정

int n, m; 
int p[MAX_N][LOG]; // 부모 정보를 저장할 배열
int depth[MAX_N]; // 깊이 정보를 저장할 배열
long long cost[MAX_N]; // 비용 정보를 저장할 배열
int v[MAX_N]; // dfs시 방문 처리를 사용할 배열
vector<pair<int, int>> tree[MAX_N]; // 인접 리스트를 저장할 배열

void input() {
	cin >> n;
	for (int i = 0; i < n - 1; i++) { // n - 1개의 간선을 양방향 인접리스트로 저장해 준다.
		int a, b, c; cin >> a >> b >> c;
		tree[a].push_back({ b, c });
		tree[b].push_back({ a, c });
	}
}

void dfs(int node, int par, int dep, long long cos) { // 루트 노드로 부터 각 노드의 깊이, 비용과 자신의 첫번째 부모 탐색
	p[node][0] = par;
	depth[node] = dep;
	cost[node] = cos;
	
	for (pair<int, int> next : tree[node]) {
		int child = next.first, c = next.second;
		if (!v[child]) {
			v[child] = 1;
			dfs(child, node, dep + 1, cos + c);
		}
	}
}

void pre_process() { // 자신의 2^j번째 부모를 찾아 저장한다.
	for (int j = 1; j < LOG; j++) {
		for (int i = 1; i <= n; i++) {
			if (p[i][j - 1] != -1) p[i][j] = p[p[i][j - 1]][j - 1];
		}
	}
}

int lca(int a, int b) { // 최소 공통 부모를 찾아준다.
	if (depth[a] < depth[b]) swap(a, b); // 비교하기 쉽게 a가 항상 크도록 유지
	
	int diff = depth[a] - depth[b]; // a와 b의 깊이 차이
	for (int i = 0; i < LOG; i++) { // a와 b의 깊이를 똑같이 맞춰주기
		if ((diff >> i) & 1) a = p[a][i];
	}

	if (a == b) return a; // b가 a의 부모인 경우

	for (int i = LOG - 1; i >= 0; i--) { // 최소 공통 부모의 바로 밑까지 위치하기
		if (p[a][i] != p[b][i]) { // 둘의 부모가 다르면 2^i만큼 깊이를 올려준다.
			a = p[a][i];
			b = p[b][i];
		}
	}

	return p[a][0]; // 둘의 바로 위 노드가 LCA임
}

int Find(int a, int d) { // a는 현재 노드, d는 올라가야할 깊이
	for (int i = 0; i < LOG; i++) {
		if ((d >> i) & 1) a = p[a][i]; // d의 i번째 비트가 존재하면 a는 2^i 번째 상위 노드로 변경
	}
	return a;
}

void query() {
	cin >> m;
	while (m--) {
		int a; cin >> a;
		if (a == 1) { // 비용 구하는 case
			int n1, n2; cin >> n1 >> n2;
			int par = lca(n1, n2); // 두 노드의 LCA 구하기
			cout << cost[n1] + cost[n2] - cost[par] * 2 << "\n"; // 루트~n1 거리 + 루트~n2 거리 - (루트~par 거리 * 2)
		}
		else {
			int n1, n2, k; cin >> n1 >> n2 >> k;
			k--; // 예제를 돌려보니 탐색 시 자기 자신을 포함한다, k는 1부터 시작이니 -1 처리
			int par = lca(n1, n2); // 두 노드의 LCA 구하기
			int left = depth[n1] - depth[par]; // n1에서 par까지의 깊이 차이
			int right = depth[n2] - depth[par]; // n2에서 par까지의 깊이 차이
			if (left >= k) cout << Find(n1, k) << "\n"; // 왼쪽에서만 탐색 가능한 경우
			else cout << Find(n2, right - (k - left)) << "\n"; // 오른쪽에서 탐색해 줘야 하는 경우
		}
	}
}

void solution() {
	v[1] = 1;
	dfs(1, -1, 0, 0);
	pre_process();
}

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

	input();
	solution();
	query();
}