백준1761 정점들의 거리 - C++

대표적인 LCA(최소 공통 조상, Lowest Common Ancestor) 문제이다.

 

정점을 트리로 만들고 두 노드 사이의 거리를 구하는 건 간단하다.

한 정점에서 출발해서 트리를 타고 가장 단거리로 다른 노드에 

도착하기만 하면 되는 문제이기 때문이다.

 

하지만 이 문제는 2초의 시간 제한이 있고,

1만 개의 질문이 주어지므로 

한 질문 당 20000 번 이하의 연산을 해야 한다.

(왜냐하면 1만 x 2만 = 2억이기 때문에 C++ 에서 러프하게 2초로 잡을 수 있다.)

 

노드가 최대 4만개이기 때문에, 재수 없게도 

일자로 쭉 펼쳐진 ㅡㅡㅡㅡㅡㅡㅡ 이런 트리일 때 양 끝의 노드만 계속 묻는다면

반드시 시간 초과가 뜬다.

물론 이런 극한의 케이스가 아니어도 보통은 시간 초과가 뜬다.

평균적으로 떨어져있는 기대값이 2만이기 때문이다.

 

고로 트리 탐색을 O(N) 수준이 아닌 O(log N) 수준으로

만들어야 할 필요가 있고, 여기서 쓰이는 것이 LCA이다.

 

LCA는 각 노드마다

자신의 2^n 번째 조상을 모두 기억하고 있게 만들고,

조상이 같지 않은 구간을 빠르게 뛰어넘은 후

처음으로 조상이 겹치는 노드를 찾아서

계산하는 원리이다.

 

#include <cstdio>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;

int N, M;
int a, b, c;
vector<pair<int, int>> edge[40001];//node , cost
int parent[40001][16];
int cost[40001];
bool visited[40001];
int depth[40001];

//Tree 생성
void makeTree() {
	queue<pair<pair<int, int>, int>> q;// node, cost, depth
	visited[1] = 1;
	q.push({ make_pair(1, 0), 0 });
	while (!q.empty()) {
		int cur_node = q.front().first.first;
		int cur_cost = q.front().first.second;
		int cur_depth = q.front().second;
		q.pop();

		for (int i = 0; i < edge[cur_node].size(); i++) {
			int next_node = edge[cur_node][i].first;
			int next_cost = edge[cur_node][i].second;
			int next_depth = cur_depth + 1;
			if (visited[next_node]) continue;
			visited[next_node] = true;
			depth[next_node] = next_depth;
			parent[next_node][0] = cur_node;
            //위에서 내려갈 때 필요한 누적 값을 저장
			cost[next_node] = cur_cost + next_cost;
			q.push({ make_pair(next_node, cur_cost + next_cost),next_depth });

		}
	}
	//LCA 배열 만들기
	for (int i = 1; i <16; i++) {
		for (int j = 1; j <= N; j++) {
			parent[j][i] = parent[parent[j][i - 1]][i - 1];
		}
	}
}

int query(int start, int end) {
	int ret = 0;
	if (depth[start] < depth[end]) {
		swap(start, end);
	}
	int depth_s = depth[start];
	int depth_e = depth[end];
	//트리의 깊이를 같은 깊이로 맞추기
	for (int i = 15; i >= 0; i--) {
		if (depth_s - (1 << i) >= depth_e) {
			depth_s -= (1 << i);
			ret += cost[start];
			start = parent[start][i];
			ret -= cost[start];
		}
	}
    //같은 깊이로 맞췄더니 만난 경우
	if (start == end) return ret;
	//가장 먼 조상부터 공통 조상 찾기
	for (int i = 15; i >= 0; i--) {
    	//조상이 다를 경우에만 끌어 올리기 
		if (parent[start][i] != parent[end][i]) {
        	//1번 조상에서 내려갈 때 드는 누적 비용 = cost[node] 
			ret += cost[start] + cost[end];
			start = parent[start][i];
			end = parent[end][i];
            //1번 조상에서 끌어 올린 곳까지 가는 비용을 빼줘야
            //start, end 노드에서 현재 위치까지 올라오는 비용을 알 수 있다.
			ret -= (cost[start]+cost[end]);
		}
	}
	//최초로 만나는 조상에서 거리
	ret += cost[start] + cost[end];
	start = parent[start][0];
    //최초 조상이 1번 조상에서 떨어진 거리 빼주기 (누적합이므로)
	ret -= (cost[start] *2);
	return ret;

}

int main() {
	scanf("%d", &N);
	for (int i = 0; i < N - 1; i++) {
		scanf("%d %d %d", &a, &b, &c);
		edge[a].push_back(make_pair(b, c));
		edge[b].push_back(make_pair(a, c));

	}
	makeTree();

	scanf("%d", &M);
	for (int i = 0; i < M; i++) {
		scanf("%d %d", &a, &b);
		printf("%d\n", query(a, b));
	}
}