[Baekjoon] 1761번: 정점들의 거리 (Platinum) - C++ 풀이

트리 최소 공통 조상 희소 배열

💾 메모리: 9268 KB, 시간: 24 ms📅 2026년 2월 25일 17:02:31

문제 링크

문제 설명

N(2 ≤ N ≤ 40,000)개의 정점으로 이루어진 트리가 주어지고 M(1 ≤ M ≤ 10,000)개의 두 노드 쌍을 입력받을 때 두 노드 사이의 거리를 출력하라.

입력

첫째 줄에 노드의 개수 N이 입력되고 다음 N-1개의 줄에 트리 상에 연결된 두 점과 거리를 입력받는다. 그 다음 줄에 M이 주어지고, 다음 M개의 줄에 거리를 알고 싶은 노드 쌍이 한 줄에 한 쌍씩 입력된다. 두 점 사이의 거리는 10,000보다 작거나 같은 자연수이다.

정점은 1번부터 N번까지 번호가 매겨져 있다.

출력

M개의 줄에 차례대로 입력받은 두 노드 사이의 거리를 출력한다.

코드 (C++)

#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <vector>

using namespace std;

const int MAX_N = 40'001;
const int MAX_M = 10'001;

int parent[MAX_N];
bool visited[MAX_N];
vector<pair<int, int>> adj[MAX_N];
int dist_from_root[MAX_N]; //루트 노드에서부터 각 노드까지의 실제 간선 가중치 누적 합을 저장
vector<pair<int, int>> queries[MAX_N];
int answers[MAX_M];


int find(int tar) {
	if (tar == parent[tar]) return tar;
	return parent[tar] = find(parent[tar]);
}

void LCA(int curr, int cost) {
	visited[curr] = true;
	dist_from_root[curr] = cost; // 루트로부터의 누적 거리 저장

	for (auto& edge : adj[curr]) {
		int next = edge.first;
		int weight = edge.second;

		if (!visited[next]) {
			LCA(next, cost + weight);
			// set union
			parent[next] = curr;
		}
		
	}

	for (auto& query : queries[curr]) {
		int tar = query.first;
		int query_idx = query.second;

		if (visited[tar]) {
			int lca_node = find(tar);
			// 공통 조상 lca_node를 기준으로 curr , tar는 연결되어있음.
			// (1~curr의 거리 - 1~lca_node) = curr~lca_node의 길이
			// (1~tar의 거리 - 1~lca_node) = tar~lca_node의 길이
			// curr~lca_node + lca_node~tar == curr~tar
			answers[query_idx] = (
				dist_from_root[curr] 
				+ dist_from_root[tar] 
				- 2 * dist_from_root[lca_node]
			);
		}
	}
	
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	//freopen("test.txt", "r", stdin);
	int u, v, w;
	int N, M;
	if (!(cin >> N)) return 0;

	for (int i = 1; i <= N; i++) parent[i] = i;

	for (int i = 0; i < N - 1; i++) {
		cin >> u >> v >> w;
		adj[u].push_back({ v, w });
		adj[v].push_back({ u, w });
	}

	if (!(cin >> M)) return 0;
	for (int i = 0; i < M; i++) {
		cin >> u >> v;
		queries[u].push_back({ v, i });
		queries[v].push_back({ u, i });
	}

	LCA(1,0);

	for (int i = 0; i < M; i++) {
		cout << answers[i] << "\n";
	}

	return 0;
}