백준 11438 LCA 2 C++

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

11438번: LCA 2

(회사에서 프로그래밍 시험을 보는데 C/C++, Java 만 있어서 C++로 다시 연습 중이다.)

 

(처음에는 자식 노드한테 번호를 순차 부여해서

- ex. (1, 2), (1, 3), 연결 시, 1 -> "1", 2 -> "10" , 3 -> "11", (2, 4), (2,5) 연결 시, 4 -> "100", 5 -> "101", .... 

prefix 일치 여부로 풀려고 했으나 끝없는 메모리 초과.... 그렇게 안 잡아먹을것 같은데...???)

 

LCA를 구하는 문제입니다. 하지만 N의 최대값이 100,000 이고, M (노드 쌍) 최대값이 100,000 이므로, 한 칸 씩 부모 노드를 올라가는 기본 방식 사용 시 시간초과가 뜹니다. (한 칸 씩 올라가기 : O(n) => 10만 * 10만 번 해야함)

 

그렇기 때문에 dp 응용버전인 2^i 번째 부모 노드를 기준으로 탐색을 진행하면 됩니다. 방법은 다음과 같습니다.

 

1. 트리의 각 노드에 대해 depth와 2^i 번째 부모를 구한다. 2^i 번째 부모는 2^(i-1) 번쨰 부모의 2^(i-1) 번째 부모이다.

 

2. 두 노드의 LCA를 구한다. 

- 먼저 두 노드의 부모를 대상으로 depth가 같아질때까지 계산한다. (기준 : 더 depth 큰 노드, depth 작은 노드는 가만히)

-- depth 차이를 기준으로 차이의 binary 값을 이용해 부모 노드를 업데이트한다. 

--- ex. 5 = 101 => 2^0 으로 업데이트, 업데이트된 노드의 2^2 번째 부모로 업데이트 = 원래 노드의 5번째 부모

 

3. depth 같아진 노드를 대상으로 2^i 번째 부모를 대상으로 올라간다. (기준 : 두 노드 모두)

- 이 때 일치하는 2^i 번째 부모가 있더라도, 더 짧은 위치의 동일 부모가 존재할 수 있으므로 대기한다. (ex. 2^2 번째 부모 일치하더라도, 2^2 + 2^0 번째, ..., 2^0 번쨰 부모 일치 가능)

-- 일치하지 않는 경우, 해당 부모로 노드를 업데이트한다.

 

4. 최종적으로 업데이트된 두 노드의 부모 노드가 정답이다.

 

코드는 다음과 같습니다.

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

using namespace std;

// parents[node][i] : node의 2^i 번째 부모 
// 2^i 번쨰 부모 계산 : dp 이용 
// depths[node] : node 의 깊이

int n, m;
vector<int> tree[100001];
int parents[100001][20] = { {0, }, };
int depths[100001] = {0, };
int result = 0;

void init()
{
	cin >> n;
	int a, b;
	for(int i = 0; i < (n-1); i++) {
		cin >> a >> b;
		tree[a].push_back(b);
		tree[b].push_back(a);
	}
	
	return;
}

void make_tree()
{
	int node, next_node;
	int height = 0;
	int brothers = 0;
	
	queue<int> q;
	q.push(1);
	parents[1][0] = 1; // 1 부모 1 
	 
	while(!q.empty()) {
		if(brothers == 0) {
			// 부모의 형제 노드 탐색 완료 
			brothers = q.size();
			height++;
		}
		
		node = q.front();
		q.pop(); 
		
		for(auto v : tree[node]) {
			next_node = v;
			
			// 부모 이미 아는 경우  
			if(parents[next_node][0] != 0)
				continue;
				
			parents[next_node][0] = node;
			depths[next_node] = height;
			q.push(next_node);
		}
		
		brothers--;
	}
	
	return;
}

void calc_ith_parents()
{
	for(int i = 1; i < 20; i++) {
		for(int node = 2; node <= n; node++) {
			// node 의 2^i 번쨰 부모 = 2^(i-1) 번쨰 부모의 2^(i-1) 번쨰 부모  
			parents[node][i] = parents[ parents[node][i-1] ][i-1];
		}
	}
	
	return;
}

void game(int node1, int node2)
{
	// node1, node2 대상으로 2^i 씩 올라가면서 부모 depth 맞추기 
	// ex. depth 3 차이 시, 2^1 후 2^0 번째 부모에서 일치 
	// 이후 2^i 씩 depth 올리기 
	
	// 둘 중 하나 1인 경우 
	if(node1 == 1 || node2 == 1) {
		result = 1;
		return;
	}
	
	// 노드 일치 
	if(node1 == node2) {
		result = node1;
		return;
	} 
	
	// 둘 중 하나가 부모 
	if(node1 == parents[node2][0]) {
		result = node1;
		return;
	}
	
	if(node2 == parents[node1][0]) {
		result = node2;
		return;
	}
	
	// 부모 일치 
	if(parents[node1][0] == parents[node2][0]) {
		result = parents[node1][0];
		return;
	}
	
	// 기준 : node1 이 더 아래 
	if(depths[node1] < depths[node2]) {
		swap(node1, node2);
	}
	
	// depth 맞추기 
	if(depths[node1] != depths[node2]) {
		int diff = depths[node1] - depths[node2];
		int cnt = 0;
		while(diff > 0) {
			if(diff & 1) 
				node1 = parents[node1][cnt];
				
			cnt++;
			diff >>= 1;
		}
	}
	
	// 찾아진 경우 
	if(node1 == node2) {
		result = node1;
		return;
	} 
	
	// 2^i 씩 올려보기 
	for(int i = 19; i >= 0; i--) {
		// 트리 넘는 경우 
		if(parents[node1][i] == 0 || parents[node2][i] == 0)
			continue;
		
		// 일치하는 경우, 더 짧은 부모 있을 수 있음 
		if(parents[node1][i] == parents[node2][i]) 
			continue;
		
		node1 = parents[node1][i];
		node2 = parents[node2][i];
	}
	
	result = parents[node1][0];
	return;
}

int main()
{
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	
	//freopen("input.txt", "r", stdin);
	//freopen("output.txt", "w", stdout);
	
	init();
	make_tree();
	calc_ith_parents();
	
	cin >> m;
	int node1, node2;
	while(m--) {
		cin >> node1 >> node2;
		game(node1, node2);
		cout << result << '\n';
	}
}