백준 7812: 중앙 트리

https://www.acmicpc.net/problem/7812

7812번: 중앙 트리

 

풀이

$O(N^2)$으로는 TLE를 받게 됩니다. 따라서 이보다 더 빠른 코드를 찾아야 하고, 트리 dp를 생각해봤을 때 $O(N)$의 풀이를 찾을 수 있습니다.

문제의 예제 입력 테스트 케이스 1번을 봅시다.

편의를 위해 0번 노드를 루트 노드로 잡겠습니다.

$dp[i]=(i$번 노드에서 모든 정점으로 이르는 비용의 합$)$이라 합시다. 이 식을 가지고 자식 노드와 부모 노드 간의 점화식을 세우겠습니다.

 

위 그림에서 예시를 들어 $dp[1]$과 $dp[3]$의 관계식을 세워보겠습니다. $i$번 노드를 루트로 하는 서브 트리의 노드 개수를 $sub[i]$라 합니다. 그러면 $sub[1]=4$, $sub[3]=2$가 됩니다.

전체 트리에서 3번 노드가 루트인 서브 트리를 제외한 트리를 생각해봅시다. 그 트리 안에 1번 노드가 포함되어 있고, 1번 노드에서 다른 노드로 이르는 거리의 합을 $x$라 잡습니다.

3번 노드가 루트인 서브 트리에서 3번 노드에서 다른 노드로 이르는 거리의 합을 $y$라 잡습니다.

그리고 1번 노드와 3번 노드 사이의 거리 7을 $dis$라고 합시다.

 

그렇다면 $dp[1]=x+y+sub[3]\times dis$로 표현할 수 있게 됩니다(왜인지는 너무 자명하여 설명하지 않겠습니다). 그리고 $dp[3]=x+y+(($전체 노드의 개수$)-sub[3])\times dis$로 표현이 됩니다. 

위의 두 식으로 $dp[1]$과 $dp[3]$과의 관계식을 구하면 $dp[3]=dp[1]-(2sub[3]-N)\times dis$(단, $N$은 전체 노드의 개수)가 됩니다.

 

예시와 같이 자식과 부모 간의 점화식을 구하면 $dp[parent]=dp[subling]-(2sub[subling]-N)\times dis$가 됩니다.

초기조건인 루트의 dp값은 처음 dfs를 돌리면서 구해줍니다.

 

코드

위의 풀이로 짠 코드입니다.

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
vector<pair<ll, ll>> adj[10000];
bool visited[10000];
ll sub[10000], dp[10000], n;
ll subdfs(ll x) {
    ll s=1, tot=0;
    for(auto i: adj[x]) {
        if(!visited[i.first]) {
            visited[i.first]=1;
            tot+=subdfs(i.first);
            s+=sub[i.first];
            tot+=i.second*sub[i.first];
        }
    }
    sub[x]=s;
    return tot;
}
void dfs(ll x) {
    for(auto i: adj[x]) {
        if(!visited[i.first]) {
            visited[i.first]=1;
            dp[i.first]=dp[x]-(2*sub[i.first]-n)*i.second;
            dfs(i.first);
        }
    }
}
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    ll a, b, w;
    cin >> n;
    while(n) {
        for(ll i=0; i<n; i++) {
            adj[i].clear();
            visited[i]=0;
            sub[i]=0;
            dp[i]=0;
        }
        for(ll i=0; i<n-1; i++) {
            cin >> a >> b >> w;
            adj[a].push_back({b, w});
            adj[b].push_back({a, w});
        }
        visited[0]=1;
        dp[0]=subdfs(0);
        for(ll i=0; i<n; i++) visited[i]=0;
        visited[0]=1;
        dfs(0);
        ll m=LONG_MAX;
        for(ll i=0; i<n; i++) m=min(m, dp[i]);
        cout << m << "\n";
        cin >> n;
    }
}