본문 바로가기

알고리즘 문제풀이/백준 문제풀이

[BOJ 7812] 중앙 트리

반응형

 

[문제]

 

www.acmicpc.net/problem/7812

 

7812번: 중앙 트리

입력은 여러 개의 테스트 케이스로 이루어져 있다. 각 테스트 케이스의 첫 줄에는 트리의 정점의 수 n이 주어진다. (1 ≤ n ≤ 10,000) 각 정점은 0번부터 n-1번까지 번호가 붙여져 있다. 다음 n-1개 줄

www.acmicpc.net

 

[난이도]

 

- Platinum III (20.12.14 기준)

 

 

[필요 개념]

 

- Tree DP 

- DFS

 

 

[풀이]

 

설명에서 나오는 S(u)는 'u를 루트로 하는 서브트리'를 뜻한다.

 

한 점에 대해서 다른 모든 점까지의 거리의 합을 구한다면 단순 Tree DP를 이용하여 노드의 개수만큼 탐색하면 된다. 

하지만, 이 문제는 모든 정점에 대해서 고려를 해야하기 때문에 $O(n^2)$이고, 여러 테스트 케이스가 들어오기 때문에 시간 내에 통과하지 못한다. 

따라서 $O(n)$안에 각 정점별로 모든 정점에 이르는 비용의 합을 구할 수 있어야 한다. 

 

우선, 각 정점별로 자신을 루트로 하는 서브트리의 노드 수를 저장해 두고,

전체 트리의 루트를 1이라고 하고 정점 1로부터 모든 점까지의 거리의 합을 구해둔다. 

그다음, DFS 순회로 트리를 타고 내려가면서 해당 노드로부터 모든 점까지의 거리의 합을 구한다.

dp[u] = 'u로부터 모든 점까지의 거리의 합'이라고 하고, 현재 노드를 v, 부모 노드를 u라고 할 때,

dp[u]를 이용해서 dp[v]를 구하는 과정을 알아보자. 

 

먼저, dp[u]는 위 그림을 참고하면 (1) + (2) 와 같다.

(1) S(v)의 모든 노드로부터 u까지의 거리의 합

(2) N-S(v)의 모든 노드로부터 u까지의 거리의 합

 

(1)에서 S(v)의 노드 수 * (u, v)의 가중치를 빼면, S(v)의 모든 노드로부터 v까지의 거리의 합과 동일하다.

(2)에서는, 원래 N-S(v)의 노드들이 u까지만 도달했기 때문에, v까지 도달하기 위해서는 N-S(v)의 노드 수 * (u, v) 만큼의 거리가 더해져야 한다. 

 

따라서 결국 dp[v]는 

dp[v] = dp[u] + (N - S(v)의 노드 수*2) * (u, v)의 가중치이다.

 

 

[소스 코드]

 

#include<bits/stdc++.h>
#define ini(x, y) memset(x, y, sizeof(x));
using namespace std;
using ll = long long;
using pii = pair<int, int>;
vector<vector<pii>>adj;
ll dis[10001];
int sub[10001];
int n;
ll init(int u, int p) {
	sub[u] = 1;
	dis[u] = 0;
	for (pii v : adj[u]) {
		if (v.first == p) continue;
		dis[u] += init(v.first, u);
		dis[u] += v.second*sub[v.first];
		sub[u] += sub[v.first];
	}
	return dis[u];
}
void sol(int u, int p) {
	for (pii v : adj[u]) {
		if (v.first == p) continue;
		dis[v.first] = dis[u] + (n - 2 * sub[v.first])*v.second;
		sol(v.first, u);
	}
}
int main(void) {
	ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
	while (1) {
		cin >> n;
		if (n == 0) break;
		ini(dis, 0);
		adj.assign(n + 1, vector<pii>());
		for (int i = 1; i < n; i++) {
			int a, b, c; cin >> a >> b >> c;
			a++; b++;
			adj[a].push_back({ b, c });
			adj[b].push_back({ a, c });
		}
		init(1, 1);
		sol(1, 1);
		ll ans =1e18;
		for (int i = 1; i <= n; i++) ans = min(ans, dis[i]);
		cout << ans << '\n';
	}
}

 

PC로 보시는 것을 권장합니다. 

피드백은 언제나 환영입니다. 댓글로 달아주세요 ^-^

반응형