【换根DP】CF1882 D-蒲公英云

【换根DP】CF1882 D

Bertha 。 2024-02-21 03:15 203阅读 0赞

Problem - D - Codeforces

思路：

一个很套路的换根

首先观察到，先对儿子一定比先对父亲操作来的代价小，因此考虑先对儿子操作，再对父亲操作

然后就可以直接换根了，首先考虑树形DP，设dp[u] 为把 u 子树染成同一种颜色的最小代价

那么根据刚刚的先操作儿子，转移方程为

dp[u] += dp[v] + (a[u] ^ a[v]) * sz[v]

然后换根就正常换就好了

Code：

#include <bits/stdc++.h>
#define int long long
constexpr int N = 2e5 + 10;
constexpr int mod = 998244353;
std::vector<int> adj[N];
int n;
int a[N];
int dp[N], sz[N];
void dfs1(int u, int fa) {
    sz[u] = 1;
    for (auto v : adj[u]) {
        if (v == fa) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        dp[u] += dp[v] + (a[u] ^ a[v]) * sz[v];
    }
}
void dfs2(int u, int fa) {
    for (auto v : adj[u]) {
        if (v == fa) continue;
        dp[u] -= (a[u] ^ a[v]) * sz[v] + dp[v];
        sz[u] -= sz[v];
        int t1 = (a[u] ^ a[v]) * sz[v] + dp[v];
        int t2 = sz[v];
        dp[v] += (a[v] ^ a[u]) * sz[u] + dp[u];
        sz[v] += sz[u];
        dfs2(v, u);
        dp[u] += t1;
        sz[u] += t2;
    }
}
void solve() {
    std::cin >> n;
    for (int i = 1; i <= n; i ++) {
        adj[i].clear();
        dp[i] = sz[i] = 0;
    }
    for (int i = 1; i <= n; i ++) {
        std::cin >> a[i];
    }
    for (int i = 1; i <= n - 1; i ++) {
        int u, v;
        std::cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 0);
    for (int i = 1; i <= n; i ++) {
        std::cout << dp[i] << " \n" [i == n];
    }
}
signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int t = 1;
    std::cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}