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

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

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

樣例輸入
復制樣例數(shù)據(jù)
3
1 0 1
1 2 1
1 3 2
樣例輸出
8

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

#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的個數(shù)dp[u][0]+=dp[to][0];//記錄0的個數(shù)} } 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（中国石油大学我要变强第九场）的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。