【换根DP】CF1882 D

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

[Problem - D - Codeforces][]

![deb798f4239b42af8fc9c94f32fdb2aa.png][]

思路：

一个很套路的换根

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

然后就可以直接换根了，首先考虑树形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;
    }

[Problem - D - Codeforces]: https://codeforces.com/contest/1882/problem/D
[deb798f4239b42af8fc9c94f32fdb2aa.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/02/21/aa1f63c8f7d444f493bf6bf36f64ed91.png