백준 3584번 가장 가까운 공통 조상 C++ 최소 공통 조상, 유니온 파인드

리뷰

 

 

문제 풀이

1. 테스트 케이스를 입력 받고, 해당 테스트 케이스 만큼 while문을 돌려준다.
2. 각 테스트 케이스 마다 n, root, tree, parent정보를 초기화 해준다.
3. 유니온 파인드용 노드 배열을 노드의 1 ~ n 노드까지 자기 자신을 값으로 갖게끔 인덱스를 설정해 준다.
4. 이후 n - 1개의 간선을 입력 받고, tree벡터 배열에 양방향 간선으로 추가, b를 a의 자식으로 유니온 해준다.
5. 1부터 n까지 노드를 탐색해 자기 자신을 value로 갖고있는 노드를 찾고 root로 할당해 준다.
6. dfs함수를를 root 부터 돌려 각 자신의 바로 위 부모를 구해준다.
7. preprocess 함수를 통해 그 위의 부모들까지 초기화를 해준다.
8. 비교할 노드 두개를 입력 받고 lca 함수의 매개변수로 전달해준 뒤 리턴값을 출력해 주면 된다.

 

참고 사항

루트 노드가 항상 1이 아니고 매번 바뀌기 때문에 유니온 파인드를 통해 루트 로드를 구해주었다.

계속 틀리다가 parent 배열을 init 함수에서 초기화 해주니 AC를 받았다.

 

정답 코드

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

using namespace std;

const int MAX_N = 10001;
const int LOG = 15;
int parent[MAX_N][LOG];
vector<int> tree[MAX_N];
int nodes[MAX_N];
int depth[MAX_N];
int tc, n, root;

void init() { // n, root, tree, parent 초기화
	cin >> n;
	root = -1;
	for (int i = 0; i < MAX_N; i++) tree[i].clear();
	memset(parent, 0, sizeof(parent));
}

int Find(int a) {
	if (nodes[a] == a) return a;
	return nodes[a] = Find(nodes[a]);
}

void Union(int a, int b) {
	int A = Find(a);
	int B = Find(b);
	if (A == B) return;
	nodes[B] = A;
}

void dfs(int node, int par, int dep) { // 각 노드의 바로 위 부모 노드 초기화(연결리스트 활용)
	parent[node][0] = par;
	depth[node] = dep;
	for (int child : tree[node]) {
		if (child != par) dfs(child, node, dep + 1);
	}
}

void preprocess() { // 각 노드의 모든 부모 초기화 (없을 경우 -1)
	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]; // 깊이 차이 계산

	for (int i = 0; i < LOG; i++) { // 깊이 차이만큼 u를 상위 노드로 올려준다.
		if ((diff >> i) & 1) u = parent[u][i];
	}

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

	for (int i = LOG - 1; i >= 0; i--) {
		if (parent[u][i] != parent[v][i]) { // 부모가 다를 경우 깊이를 2^i씩 높여준다.
			u = parent[u][i];
			v = parent[v][i];
		}
	}

	return parent[u][0]; // 노드의 바로 위 부모 노드를 반환
}

int main() {
	cin >> tc;
	while (tc--) {
		init(); // 초기화
		for (int i = 1; i <= n; i++) nodes[i] = i; // 유니온 파인드용 노드 배열 초기화
		for (int i = 0; i < n - 1; i++) { // 인접리스트 받아오기 & 유니온 처리하기
			int a, b; cin >> a >> b;
			tree[a].push_back(b);
			tree[b].push_back(a);
			Union(a, b);
		}
		for (int i = 1; i <= n; i++) { // 루트 노드를 찾는다.
			//cout << i << " " << nodes[i] << "\n";
			if (nodes[i] == i) {
				//cout << "ROOT : " << i << "\n";
				root = i;
				break;
			}
		}
		int n1, n2; cin >> n1 >> n2;
		dfs(root, -1, 0); // 루트를 기준으로 dfs 실행
		preprocess();
		cout << lca(n1, n2) << "\n";
	}
}