백준 20188 : 등산 마니아 (트리 DP)

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

20188번: 등산 마니아

풀이
1. 완전탐색 + 최소공통조상

한 경로의 다양성은 루트로부터 i와 j의 거리를 합하고 루트와 i와 j의 공통 조상의 거리를 뺀 값과 같다.
모든 i와 j에 대해 다양성을 구하는 시간복잡도는 n^2 logn, n이 30만이므로 시간초과

 

2. 트리 DP
하나의 서브트리에 대해
1. 자식노드와 루트를 골랐을 때
2. 자식노드와 다른 방향 자식노드를 골랐을 때
두 가지로 나눈다.
첫번째 경우는 자식노드의 깊이와 같으므로 자식노드들의 깊이의 합이된다.
두번째 경우는 자식노드의 깊이의 합 + 다른쪽 자식노드의 깊이의 합이된다.
이때 경우의 수가 여러개 나올 수 있으므로 모든 경우에 대해 합해준다.

#include<bits/stdc++.h>
#include<time.h>


#define MAX 1000000007
#define INF 2000000000
#define MOD 1000000007 
using namespace std;
typedef long long ll;
typedef pair<int, int > pii;
typedef pair<long long, long long> pll;
typedef pair<double, double> pdd;

int dx[8] = { 0, -1, 0, 1, 1, 1, -1, -1 };
int dy[8] = { -1, 0, 1, 0, -1, 1, 1, -1 };

/*
int dx[2][6] = {
    {1, -1, 0, 0 ,1, -1},
    {1, -1, 0, 0, -1, 1}
};
int dy[2][6] = {
    {0, 0, 1, -1 ,1, -1},
    {0, 0, 1, -1, 1, -1}
};*/

//int dx[8] = { -1, -1, 1, 1, -2, -2, 2, 2 };
//int dy[8] = { 2, -2, 2, -2, -1, 1, -1, 1 };

void init() {
    ios::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
}
ll gcd(ll a, ll b) { for (; b; a %= b, swap(a, b)); return a; }
ll lcm(ll a, ll b) { return (a * b) / gcd(a, b); }

vector<int> graph[300001];
vector<int> tree[300001];
pll dp[300001];

ll deg[300001];
int n;
ll ans = 0;
void Bfs() {
    queue<int> q;
    q.push(1);
    deg[1] = 0;
    while (!q.empty()) {
        int x = q.front();
        q.pop();
        for (int next : graph[x]) {
            if (deg[next] == -1) {
                //parent[next][0] = x;
                tree[x].push_back(next);
                deg[next] = deg[x] + 1;
                q.push(next);
            }
        }
    }
}

pll func(int x) {
    if (tree[x].size() == 0) {
        return dp[x]= { deg[x],1 };
    }
    ll cnt = 0;
    ll val = 0;

    for (int i = 0; i < tree[x].size(); i++) {
        int child = tree[x][i];
        if (dp[child] == make_pair(0ll, 0ll)) {
            func(child);
        }
        cnt += dp[child].second;
        val += dp[child].first;
    }
    
    ll cntt = 0;
    for (int i = 0; i < tree[x].size(); i++) {
        int child = tree[x][i];
;       ans += dp[child].first * (cnt - dp[child].second);
        cntt += dp[child].second * (cnt - dp[child].second);
    }
    ans -= (cntt / 2) * deg[x];
    ans += val;
    
    return dp[x] = make_pair(val + deg[x], cnt + 1);
}

int main() {
    init();
    memset(deg, -1, sizeof(deg));
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int a, b; cin >> a >> b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    Bfs();
    func(1);
    cout << ans << "\n";
}