백준20188 등산 마니아

작성일 | In

문제 링크

  • http://icpc.me/20188

문제 출처

  • 2020 KOI 2차 초등부 3번
  • 2020 KOI 2차 중등부 2번

사용 알고리즘

  • 트리 DP

시간복잡도

  • $O(N)$

풀이

모든 경로 길이의 합과 루트에서 어떤 두 정점의 LCA로 가는 경로로 나눠서 생각합니다.

모든 경로 길이의 합은 각 간선이 사용되는 횟수를 구해서 모두 더해주면 됩니다.
정점 $p, c$($p$가 부모 정점)를 연결하는 간선 $(p, c)$가 사용되는 횟수는 $S_c \cdot (N-S_c)$입니다.(단, $S_v$는 $v$를 루트로 하는 서브 트리의 크기)

루트에서 어떤 두 정점의 LCA로 가는 경로의 길이 합은 각 정점이 LCA가 되는 경우의 수를 구하면 됩니다. 정점 $p$의 자식들을 각각 $c_1, c_2, \cdots , c_k$라고 할 때, 정점 $p$가 LCA가 되는 경우의 수는 ${S_p \choose 2} - \sum_{i=1}^{k} {S_{c_i} \choose 2}$입니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
#include <bits/stdc++.h>
#define x first
#define y second
#define all(v) v.begin(), v.end()
#define compress(v) sort(all(v)), v.erase(unique(all(v)), v.end())
using namespace std;

typedef long long ll;

int N, S[303030], D[303030], P[303030];
vector<int> G[303030];
ll R;

ll C2(int x){ return 1LL * x * (x - 1) / 2; }

int DFS(int v, int b = -1){
    S[v] = 1; P[v] = b;
    for(const auto i : G[v]) if(i != b) D[i] = D[v] + 1, S[v] += DFS(i, v);
    return S[v];
}

int main(){
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    cin >> N;
    for(int i=1; i<N; i++){
        int s, e; cin >> s >> e;
        G[s].push_back(e); G[e].push_back(s);
    }
    DFS(1);

    for(int i=1; i<=N; i++) R += 1LL * S[i] * (N - S[i]);
    for(int i=1; i<=N; i++){
        ll now = C2(S[i]);
        for(const auto j : G[i]) if(j != P[i]) now -= C2(S[j]);
        R += now * D[i];
    }
    cout << R;
}