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

자료 구조 트리 LCA 희소 배열

문제

문제 링크

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

입력

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

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

출력

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

문제 조건

• 목표: 가중치가 있는 트리에서 주어진 두 정점 사이의 거리를 구한다.

• 입력 상태: 정점 수 N(≤40,000), 쿼리 수 M(≤10,000), 간선 가중치(≤10,000)인 트리 구조 데이터.

• 핵심 조건: 두 노드 u, v 사이의 거리는 ‘루트로부터 u까지의 거리 + 루트로부터 v까지의 거리 - 2 * (루트로부터 LCA(u, v)까지의 거리)’ 공식을 통해 계산해야 한다.

---

풀이

핵심 알고리즘

자료 구조

• 시간 복잡도: O((N + M) · α(N)) — DFS 1회 순회 및 Union-Find 연산 포함(약 50,000번 내외의 연산).

핵심 아이디어

Tarjan’s Offline LCA 알고리즘을 사용하여 모든 쿼리를 한 번의 DFS 순회로 해결한다. 각 노드의 루트로부터의 누적 거리를 저장함과 동시에, Union-Find를 이용해 실시간으로 최소 공통 조상(LCA)을 찾아 거리 공식을 적용한다.

① 오프라인 쿼리 수집 및 초기화

매 쿼리마다 트리를 탐색하면 시간 초과가 발생하므로, 쿼리 정보를 각 정점의 인접 리스트에 저장하여 DFS 순회 중 해당 정점을 만날 때 즉시 처리할 수 있도록 준비한다.

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

② 루트로부터의 누적 거리 계산

두 노드 사이의 거리를 LCA를 기점으로 분할하여 계산하기 위해, DFS 탐색 과정에서 파라미터 cost를 넘겨주며 루트(1번 노드)에서 각 정점까지의 가중치 합을 dist_from_root 배열에 기록한다.

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;
		}
		
	}

③ Union-Find를 이용한 LCA 식별 및 거리 도출

현재 노드 curr의 자식 서브트리 탐색이 완료될 때마다 부모와 합치며, 만약 쿼리 상대 노드 tar가 이미 방문된 상태라면 find(tar)를 통해 도출된 대표 원소가 두 노드의 LCA가 된다. 이를 통해 미리 계산된 누적 거리들을 공식에 대입하여 최종 거리를 구한다.

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]
			);
		}
	}

---

성능

• 메모리 : 9268 KB

• 시간 : 24 ms

코드 (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;
}

---