題目描述
“我是番茄醬！”
“我是黃芥末醬！”
“合在一起就是——美式熱狗上加的，那個！“
熱狗樹上的每個節點都涂有番茄醬或者黃芥末醬中的一種，這樣熱狗樹就變得美味了~LiMn2O4構造了一顆熱狗樹，他想知道這棵樹的美味程度。
一個熱狗樹的美味程度，定義為每個節點到其他和自己口味不一樣的節點的最短距離之和的和，也就是任意兩個口味不同的節點之間的路徑長度和。請你求出這顆樹的美味值，并且答案對998244353取模。

輸入
第一行一個正整數N。
第二行有N個數，他們的值在[0,1]范圍內取，其中1代表是涂了番茄醬，0代表是涂了黃芥末醬。
接下來N-1行，每一行有三個數a,b,w。代表節點a到b有邊，路徑的權值是w。輸入數據保證是一棵樹。
N≤100000,1≤a,b≤N,w<1 000 000 000

輸出
輸出答案對998244353取模后的結果。

樣例輸入
復制樣例數據
3
1 0 1
1 2 1
1 3 2
樣例輸出
8

提示
樹（tree）是包含n（n>=0）個結點的有窮集，其中：
（1）每個元素稱為結點（node）；
（2）有一個特定的結點被稱為根結點或樹根（root）。
（3）除根結點之外的其余數據元素被分為m（m≥0）個互不相交的集合T1，T2，……Tm-1，其中每一個集合Ti（1<=i<=m）本身也是一棵樹，被稱作原樹的子樹（subtree）。
補一下以前的坑，當時比賽的時候沒有做到這個題。。今天想起來了，就來補一下。
統計不同顏色的節點互相到達的距離和。要是我們去枚舉點到點的距離，這樣肯定是不行的，因為那樣是O（n^2） 的時間復雜度。一般這樣我么就去枚舉邊，因為邊的條數不多。那么我們去考慮一下每條邊的貢獻。每一條邊的貢獻就等于這個節點的子樹中0的個數*（1的個數-子樹中1的個數）* w+子樹中1的個數*（0的個數-子樹中0的個數）。
代碼如下：

#include<bits/stdc++.h> #define ll long long #define mod 998244353 using namespace std;const int maxx=1e5+100; struct edge{int to;int next;ll w; }e[maxx<<1]; int head[maxx],col[maxx]; ll dp[maxx][2],ans,cnt1,cnt0; int n,tot;void add(int u,int v,ll w) {e[tot].to=v;e[tot].w=w;e[tot].next=head[u];head[u]=tot++; }inline void dfs(int u,int f) {dp[u][col[u]]=1;for(int i=head[u];i!=-1;i=e[i].next){int to=e[i].to;if(to==f) continue;dfs(to,u);ans=(ans+(dp[to][1]*(cnt0-dp[to][0])%mod)*e[i].w%mod)%mod;ans=(ans+(dp[to][0]*(cnt1-dp[to][1])%mod)*e[i].w%mod)%mod;dp[u][1]+=dp[to][1];//記錄1的個數dp[u][0]+=dp[to][0];//記錄0的個數} } int main() {int l,r;ll w;while(~scanf("%d",&n)){tot=cnt1=cnt0=0;memset(head,-1,sizeof(head));for(int i=1;i<=n;i++) scanf("%d",&col[i]),cnt1+=col[i];cnt0=n-cnt1;for(int i=1;i<n;i++){scanf("%d%d%lld",&l,&r,&w);add(l,r,w);add(r,l,w);}ans=0;dfs(1,0);printf("%lld\n",ans*2%mod);//每一條邊的貢獻我們在dfs的時候只求了一遍，但是實際上一條邊該是貢獻了兩遍價值}return 0; }

努力加油a啊，(^o)/~

總結

以上是生活随笔為你收集整理的热狗树 树形dp（中国石油大学我要变强第九场）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。