[BOJ] 3176번: 도로 네트워크 (C++)

Problem

3176번: 도로 네트워크

부끄럽지만 문제를 이해하는데 시간이 조금 걸렸다. 필자처럼 문제를 이해하는데 헤매는 사람이 있을 수도 있으니 문제에 대해 정리해보겠다.

예제 입력 1을 시각화하면 다음과 같다. 이때 도시 1과 도시 2를 연결하는 경로는 2-3-1로 길이가 100인 도로와 길이가 50인 도로로 이루어져 있다. 따라서 도시 1과 도시 2를 연결하는 경로에서 가장 짧은 도로의 길이는 50, 가장 긴 도로의 길이는 100이다.

---

Solution

 

N개의 노드와 N-1개의 간선으로 이루어져 있으므로 트리 구조다.

 

트리 구조에서 노드 A와 노드 B사이의 경로를 파악하는 문제는 LCA(Least Common Ancestor) 알고리즘을 통해 각 쿼리당 O(logN) 안에 해결할 수 있다.

 

기존의 LCA 알고리즘은 length 배열에 2ⁱ번째 부모까지의 길이를 저장한다.

 

이를 변형하여 minLength 배열에는 2ⁱ번째 부모까지의 경로 중 가장 짧은 도로의 길이를 저장하고 maxLength 배열에는 2ⁱ번째 부모까지의 경로 중 가장 긴 도로의 길이를 저장하여 문제를 해결할 수 있다.

 

#include <cstdio>
#include <algorithm>
#include <vector>

using namespace std;

const int nMAX = 100001;
const int iMAX = 17;

int N;
vector<pair<int, int>> adj[nMAX];

bool visited[nMAX]; // dfs 함수를 위한 방문 여부 배열

int minLength[nMAX][iMAX]; // 최소 거리 저장
int maxLength[nMAX][iMAX]; // 최대 거리 저장
int parent[nMAX][iMAX]; // 부모 정보 저장
int depth[nMAX]; // 깊이 정보 저장

// 깊이, 부모, 거리 배열 초기화
void dfs(int start, int n){
	visited[start] = true;
	depth[start] = n;
	for(int i=0; i<adj[start].size(); i++){
		int next = adj[start][i].first;
		int weight = adj[start][i].second;
		
		if(visited[next]) continue;

		// next의 2^0번째 부모는 start
		parent[next][0] = start;
		// next에서 2^0번째 부모까지의 거리는 weight
		minLength[next][0] = weight;
		maxLength[next][0] = weight;
		dfs(next, n+1);
	}
}
// 부모, 거리 배열 Memoization
// 2^i번째 부모가 누구인지, 부모까지의 거리가 얼마인지 저장
void connect(){
	for(int p=1; p<iMAX; p++){
		for(int cur=1; cur<=N; cur++){
			int prevParent = parent[cur][p-1];
			parent[cur][p] = parent[prevParent][p-1];
			
			if(parent[prevParent][p-1] == 0) continue;
			
			int prevMin = minLength[prevParent][p-1];
			int prevMax = maxLength[prevParent][p-1];

			int curMin = minLength[cur][p-1];
			int curMax = maxLength[cur][p-1];

			minLength[cur][p] = min(prevMin, curMin);
			maxLength[cur][p] = max(prevMax, curMax);
		}
	}
}
// 최소 공통 조상을 찾으며 거리를 구한다
pair<int, int> LCA(int xNode, int yNode){
	if(depth[xNode] > depth[yNode]){
		int temp = xNode;
		xNode = yNode;
		yNode = temp;
	}
	int mins = 1000001;
	int maxs = -1;
	// 두 노드의 높이를 맞춰준다
	for(int i=iMAX-1; i>=0; i--){
		int jump = 1 << i;
		if(depth[yNode] - depth[xNode] >= jump){
			// 최대 최소 길이를 갱신한다
			mins = min(mins, minLength[yNode][i]);
			maxs = max(maxs, maxLength[yNode][i]);
			yNode = parent[yNode][i];
		}
	}
	// 두 노드가 같다면 종료한다
	if(xNode == yNode) return {mins, maxs};
	// 같지 않으면 같을 때까지 올라가며 검사한다
	for(int i=iMAX-1; i>=0; i--){
		if(parent[xNode][i] == parent[yNode][i]) continue;
		
		// 최대 최소 길이를 갱신한다
		mins = min(mins, min(minLength[xNode][i], minLength[yNode][i]));
		maxs = max(maxs, max(maxLength[xNode][i], maxLength[yNode][i]));
		
		xNode = parent[xNode][i];
		yNode = parent[yNode][i];
	}
	// 최대 최소 길이를 갱신한다
	mins = min(mins, min(minLength[xNode][0], minLength[yNode][0]));
	maxs = max(maxs, max(maxLength[xNode][0], maxLength[yNode][0]));

	return {mins, maxs};
}

int main() {
	scanf("%d", &N);
	// 트리를 만든다
	for(int i=0; i<N-1; i++) {
		int A, B, C;
		scanf("%d%d%d", &A, &B, &C);

		adj[A].push_back({B, C});
		adj[B].push_back({A, C});
	}

	int K;
	scanf("%d", &K);
	vector<pair<int, int>> ans(K);
	for(int i=0; i<K; i++) {
		int D, E;
		scanf("%d%d", &D, &E);
		ans[i] = {D, E};
	}
	// LCA를 위한 전처리
	dfs(1, 0);
	connect();

	for(int i=0; i<K; i++) {
		int x = ans[i].first;
		int y = ans[i].second;

		pair<int, int> ans = LCA(x, y);
		printf("%d %d\n", ans.first, ans.second);
	}
}

---

Result

알고리즘 상으로 시간 초과가 발생할만한 부분이 존재하지 않는 것 같은데 계속해서 시간 초과가 떠서 당황했었다. 혹시 입출력이 느려서 시간 초과가 뜨는 건가 싶어서 scanf와 printf로 바꾸니까 바로 문제가 해결되었다! 앞으로는 scanf, printf를 사용해야겠다.