[BOJ 15669] 나무 위의 입자

문제

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

15669번: 나무 위의 입자

 

---

풀이

트리DP를 사용합니다.

트리가 주어졌을 때, 두 배열을 만들어 다음과 같이 정의합니다:

 

• \(depth[i]:\) 노드 \(i\)의 깊이.
• \(dp[i][0], dp[i][1]:\) \(i\)를 루트로 하는 서브트리의 노드 중 깊이가 짝수인 노드의 개수, 홀수인 노드의 개수.

트리를 순회하면서 위 두 배열에 값을 저장한 다음 각 \(M\)개의 쿼리에 대하여 다음 작업을 수행합니다.

입력받은 \(u, v, c\)에 대하여 \(depth[u] > depth[v]\)라고 할 때:

 

• c가 0인 경우: 시작이 0, 끝이 0으로 같으므로 시작과 끝 노드의 깊이가 둘다 짝수이거나 둘 다 홀수여야 합니다.
  \(\rightarrow ans = dp[u][0](dp[1][0] - dp[u][0]) + dp[u][1](dp[1][1] - dp[u][1])\)
• c가 1인 경우: 시작이 0, 끝이 1 이므로 시작과 끝 노드의 깊이가 짝홀 또는 홀짝이어야 합니다.
  \(\rightarrow ans = dp[u][0](dp[1][1] - dp[u][1]) + dp[u][1](dp[1][0] - dp[u][0])\)

간단히 \(u\)를 루트로 하는 서브트리의 깊이가 홀 또는 짝인 노드의 개수와 전체 트리의 깊이가 홀 또는 짝인 노드의 개수에서 그 개수를 뺀 나머지 개수를 곱한 값을 구하면 됩니다.

 

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

 

---

코드

#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, M;
int depth[100001], dp[100001][2];
vector<vector<int>> graph;

void init_tree(int cur, int pnode) {
    dp[cur][depth[cur] % 2]++;
    for (auto next : graph[cur]) {
        if (next == pnode) continue;
        depth[next] = depth[cur] + 1;
        init_tree(next, cur);
        dp[cur][0] += dp[next][0];
        dp[cur][1] += dp[next][1];
    }
}

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

    cin >> N >> M;
    graph.resize(N + 1);
    for (int i = 0; i < N - 1; ++i) {
        int a, b;
        cin >> a >> b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    init_tree(1, 0);

    while (M--) {
        int u, v, c;
        ll ans;
        cin >> u >> v >> c;
        if (depth[u] < depth[v])
            swap(u, v);
        int a1 = dp[u][0], a2 = dp[u][1], b1 = dp[1][0] - a1, b2 = dp[1][1] - a2;
        if (c == 0)
            ans = 1LL * a1 * b1 + 1LL * a2 * b2;
        else
            ans = 1LL * a1 * b2 + 1LL * a2 * b1;
        cout << ans << '\n';
    }

    return 0;
}