热狗树

题面：

有一棵点权为0或1的树，问所有0的点到所有1的距离和以及所有1的点到所有0的距离和。

一个热狗树的美味程度，定义为每个节点到其他和自己口味不一样的节点的最短距离之和的和，也就是任意两个口味不同的节点之间的路径长度和。请你求出这颗树的美味值，并且答案对998244353取模。

思路：

树形dp，设num[0]和num[1]分别为点权为0或者1的点的个数，由dp可以得到u号节点的siz[0][u]和siz[1][u]分别为u节点子树点权为0和子树点权为1的点的个数，则对于某一个点u，其连着他爸爸的那条边的边权为w，则这条边对于答案的贡献为

ans+=w∗(siz[0]∗(num[1]−siz[1])+siz[1]∗(num[0]−siz[0]))

至于为什么你们自己稍微想一想应该就能理解了！（考虑一条边将一棵树分成了两个部分）记得取模。

代码：

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 typedef long long ll;
 4 const int N=100100;
 5 const int p=998244353;
 6  
 7 typedef  pair<int,int> pa;
 8 vector<pa> G[N];
 9 int a[N],si[2][N],n,s[2];
10 ll ans;
11  
12 void dfs1(int x, int fa) {
13     si[a[x]][x] = 1;
14  
15     for (auto it : G[x]) {
16         int v = it.first;
17         if (v == fa) continue;
18         dfs1(v, x);
19         si[0][x] += si[0][v];
20         si[1][x] += si[1][v];
21     }
22 }
23  
24 void dfs2(int x, int fa) {
25     for (auto it : G[x]) {
26  
27         int v = it.first;
28         int w = it.second;
29         if (v == fa) continue;
30  
31         ans = (ans+1ll * si[1][v] * (s[0] - si[0][v]) % p * w)%p;
32         ans = (ans+1ll * si[0][v] * (s[1] - si[1][v]) % p * w)%p;
33  
34         dfs2(v, x);
35     }
36 }
37  
38 int main() {
39     scanf("%d", &n);
40     for (int i = 1; i <= n; i++) {
41         scanf("%d", &a[i]);
42         s[a[i]]++;
43     }
44     for (int i = 1; i < n; i++) {
45         int u,v,w;
46         scanf("%d%d%d", &u, &v, &w);
47         G[u].emplace_back(v, w);
48         G[v].emplace_back(u, w);
49     }
50     dfs1(1, 0);
51     dfs2(1, 0);
52     printf("%lld\n", ans * 2 % p);
53 }

View Code