[G5] 백준 1240번 노드사이의 거리 C++ LCA, 트리

리뷰

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

간선에 가중치가 있는 트리로 이루어진 자료 구조에서 주어진 노드간 거리를 계산하는 문제

N, M의 최대 값이 작아 DP를 활용한 LCA를 굳이 안해도 되었지만 복기를 할 겸 해당 방법으로 시도하였다.

결국엔 반복문에서 LOG 값을 잘못 설정하여 OOB를 두번 받았다, 꾸준히 연습하는게 중요한 것 같다.

 

 

전역 변수

• LOG : 부모를 비트단위 DP로 체크하기 위한 최대값
• N : 노드 번호의 최대값
• n : 문제에 주어지는 노드의 개수, 1부터 시작
• m : 문제에 주어지는 쿼리문의 개수
• depth : 각 노드의 트리에서의 깊이를 저장할 정수형 배열
• cost : 각 노드의 루트로 부터 가중치를 저장할 정수형 배열
• parent : 각 노드의 부모 정보를 저장하기 위한 정수형 2차 배열
• Edge : 연결될 노드와 가중치를 저장하기 위한 구조체
• lst : 인접 리스트를 저장하기 위한 Edge타입 벡터 배열

 

함수

1. dfs

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

 

깊이 우선 탐색을 통해 초기 세팅을 하기 위한 함수

1. 매개 변수로 부모 노드, 현재 노드, 깊이, 가중치를 인자로 받는다.
2. parent[node][0]에 par값을 넣는다 이는 자기의 바로 위 부모를 초기화 하는 부분이다.
3. depth[node]에 dep값을 넣어 트리에서 현재 노드의 깊이 정보를 초기화 한다.
4. cost[node]에 val을 넣어 루트로 부터 현재까지의 간선의 가중치를 초기화 한다.
5. 이후 인접리스트를 순회하며 자신의 하위 노드로 뻗어나간다.
6. 자신이 부모가 되고, 자식 노드와 깊이를 증가시키고 가중치를 더해주어 재귀를 진행한다.

 

2. preprocess

void preprocess()

 

LCA를 구하기 전 부모 정보를 전처리 하기 위한 함수

1. 이미 바로 윗 부모는 dfs 함수에서 구한 상태이므로 그 이후 부모를 초기화 해준다.
2. 만약 나의 2^j-1번째 부모가 -1, 없는 상태가 아니라면 내 2^j번째 부모를 내 2^j-1번째 부모의 2^j-1번째 부모로 저장한다.
3. 아니라면 내 j번째 부모는 -1, 즉 없는 상태로 저장해 준다.

 

3. lca

int lca(int u, int v)

 

두 노드간 공통인 부모를 구하는 함수

1. 공통 부모를 구하기 위한 두 노드 u, v를 매개변수로 입력 받아준다.
2. 트리상에서 두 노드의 깊이를 비교하고 u가 더 위에 존재한다면 편의를 위해 두 노드를 바꾸어준다.
3. diff에 u의 깊이에서 v의 깊이를 뺀 값을 저장해 준다.
4. diff의 비트를 확인하며 1이 존재한다면 그 비트의 위치 만큼 u를 위로 올려준다.
5. 위 작업으로 인해 u, v의 깊이는 동일해 졌다, 만약 u와 v가 같으면 결국 u가 부모였던 경우로 바로 리턴해 준다.
6. 이후 둘의 위치를 동시에 올려주며 최종적으로 공통 부모의 바로 밑까지 이동하는 작업을 진행한다.
7. 두 노드의 2^i 번째 부모를 순회하며 공통인 부모가 아닐 경우 2^i번 만큼 두 노드를 위로 올려준다.
8. 최종적으로 u혹은 v의 바로 상위 부모가 LCA가 된다, 해당 값을 리턴해 준다.

 

문제풀이

1. n, m값을 입력 받고, n - 1개의 간선 정보를 받아 인접리스트에 양방향 간선으로 추가해 준다.
2. dfs함수를 통해 depth, cost, parent의 초기화를 진행해 준다.
3. preprocess함수를 통해 비트 단위로 부모를 탐색하여 parent배열을 완성해 준다.
4. m개의 쿼리를 입력 받고, 두 노드의 a, b를 입력 받는다.
5. a, b의 최소 공통 부모 LCA를 구한다. a와 b의 cost를 더한 값에서 LCA의 cost의 2배를 빼준 값을 출력한다.

 

참고 사항

cost는 루트로 부터 자신의 위치까지의 거리의 값이다.

따라서 두 노드 a, b의 cost를 더하고 LCA의 cost를 두 번 빼주면 두 노드간의 거리가 된다.

 

 

정답 코드

#include<iostream>
#include<vector>
#include<algorithm>
using namespace std;

const int LOG = 11;
const int N = 1001;

int n, m;
int depth[N];
int cost[N];
int parent[N][LOG];

struct Edge {
	int nn, w;
};
vector<Edge> lst[1001];

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

	for (const Edge& edge : lst[node]) {
		if (edge.nn == par) continue;
		dfs(node, edge.nn, dep + 1, val + edge.w);
	}
}

void preprocess() {
	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];
			else parent[i][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 = LOG - 1; i >= 0; --i) {
		if ((diff >> i) & 1) u = parent[u][i];
	}

	if (u == v) return 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);
	cout.tie(0);

	cin >> n >> m;
	for (int i = 0; i < n - 1; ++i) {
		int a, b, c; cin >> a >> b >> c;
		lst[a].push_back({b, c});
		lst[b].push_back({ a, c });
	}

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

	while (m--) {
		int a, b; cin >> a >> b;
		int LCA = lca(a, b);
		cout << cost[a] + cost[b] - cost[LCA] * 2 << "\n";
	}
}