[BOJ 20506] Kaisar - 생존

문제

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

20506번: Kaisar - 생존

 

---

풀이

트리DP를 사용합니다.

트리가 주어졌을 때 다음 \(2\)개의 배열을 정의합니다:

 

• \(cnt[i]:\) \(i\)를 루트로 하는 서브트리의 크기. (노드의 개수)
• \(dp[i]:\) \(i\)를 \(lca\)로 가지는 두 노드의 쌍의 개수.

\(dp[i]\), 즉, \(i\)를 \(lca\)로 가지는 두 노드의 쌍의 개수를 구하는 방법은 다음과 같습니다:

 

1. \(i\)를 루트로 하는 서브트리의 모든 쌍의 개수를 구합니다.
2. 구한 개수에서 \(i\)의 자식을 루트로 하는 서브트리에 포함된 쌍은 \(i\)를 \(lca\)로 가질 수 없으므로 그 개수를 뺍니다.

그러므로, \(i\)의 자식 노드들을 \(c_{1}, c_{2}, c_{3}, ... , c_{m}\)이라고 할 때, \(dp[i]\)에 다음 값을 저장합니다:

$$dp[i] = cnt[i] ^ 2 - \sum_{j = 1}^{m} cnt[c_{j}] ^ 2$$

이후, 구한 개수를 홀수번째의 개수와 짝수번째의 개수로 나눠 그 값을 더해나가면 됩니다.

 

시간복잡도: \(O(N)\)

 

---

코드

#include <bits/stdc++.h>
#define x first
#define y second
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

int N;
vector<vector<int>> graph;
ll cnt[200001], dp[200001], sum_a = 0, sum_b = 0;

void solve(int cur, int pnode) {
    cnt[cur] = 1;
    for (auto next : graph[cur]) {
        if (next == pnode) continue;
        solve(next, cur);
        cnt[cur] += cnt[next];
        dp[cur] -= cnt[next] * cnt[next];
    }
    dp[cur] += cnt[cur] * cnt[cur];
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int root;
    cin >> N;
    graph.resize(N + 1);
    for (int i = 1; i <= N; ++i) {
        int n;
        cin >> n;
        if (n == 0) root = i;
        graph[n].push_back(i);
    }

    solve(root, 0);

    int idx = 0, cur = 0;
    for (int i = 1; i <= N; ++i) {
        ll a = dp[i] - (dp[i] / 2), b = dp[i] / 2;
        a *= i, b *= i;
        if (!idx) {
            sum_a += a;
            sum_b += b;
        }
        else {
            sum_a += b;
            sum_b += a;
        }
        cur += dp[i];
        idx = cur % 2;
    }
    
    cout << sum_b << ' ' << sum_a << '\n';

    return 0;
}