ZR1050

ZR1050

http://www.zhengruioi.com/problem/1030

题目大意:

给定一棵带点权的树,求所有联通块的点权和的平方的和 \(n \le 10^5\)

题解

首先,关于平方的和或者和的平方我们一般都是考虑新加进来一个点会产生什么样的贡献

比如这一题,我们考虑合并两个集合的代价会发生什么样的变化

首先,设\(a,b\)分别为合并的两个集合的点权和

那么由$a^2 \(变为\)(a + b) ^ 2$的区别是代价多了个\(2ab\)和\(b^2\)

也就是说我们在DP的过程中需要维护所有联通块的点权和进行转移

接下,直接暴力计算是不可以的,我们需要对所有的联通块统一计算,发现对于合并每一对联通块的代价都是上面的东西

乘法分配律之后

我们发现还要额外记录当前点为根的联通块个数

设\(f_i\)表示以\(i\)点为根的联通块的个数

设\(g_i\)表示\(i\)点为根的所有联通块的点权和

设\(h_i\)表示以\(i\)点为根的所有联通块的点权和的平方的和

那么首先根据上面,我们在树形DP的时候可以合并两个联通块

\[h_x = (h_x + g_x\times f_y + 2\times g _x\times g_y+g_y\times f_x) \]

首先本来原有的联通块的代价肯定是要保留,另外,对于\(h_x\)可以被合并\(f_y\)次,同理\(h_y\)会被合并\(f_y\)次

最后由于\(g\)数组记录的本来就是总贡献,因此直接合并即可

同理我们可以得到\(g\)的转移

\[g_x = (g_x + g_x\times f_y +g_y\times f_x) \]

最后更新答案即可

总结

对于这种和联通块有关的题目一般可以通过单独算贡献或者DP的方式解决

特别是和平方相关的,要考虑新加进来的数的贡献

代码

#include<cstdio>
#include<iostream>
#include<queue>
#include<algorithm>
#include<cstring>
#include<cctype>
#include<vector>
#include<ctime>
#include<cmath>
#include<set>
#include<map>
#define LL long long
#define pii pair<int,int>
#define mk make_pair
#define fi first
#define se second
using namespace std;
const int N = 2e5 + 3;
const LL mod = 998244353;
int n;
LL w[N];
LL f[N],g[N],h[N];
vector <int> G[N];
inline int read(){
	int v = 0,c = 1;char ch = getchar();
	while(!isdigit(ch)){
		if(ch == '-') c = -1;
		ch = getchar();
	}
	while(isdigit(ch)){
		v = v * 10 + ch - 48;
		ch = getchar();
	}
	return v * c;
}
inline void dfs(int x,int fa){//h:平方的和 
	f[x] = 1;g[x] = w[x];
	h[x] = w[x] * w[x] % mod;
	for(int i = 0;i < (int)G[x].size();++i){
		int y = G[x][i];
		if(y == fa) continue;
		dfs(y,x);
		h[x] = (h[x] + h[x] * f[y] % mod + 2 * g[x] * g[y] % mod + h[y] * f[x] % mod) % mod;
		g[x] = (g[x] + g[x] * f[y] + g[y] * f[x]) % mod;  
		f[x] = (f[x] + f[x] * f[y]) % mod;
	}
}
int main(){
	n = read();
	for(int i = 1;i <= n;++i) w[i] = read();
	for(int i = 1;i < n;++i){
		int x = read(),y = read();
		G[x].push_back(y);
		G[y].push_back(x);	
	}
    dfs(1,0); 
    LL ans = 0;
    for(int i = 1;i <= n;++i) ans = (ans + h[i]) % mod;
    cout << ans;
	return 0;
}