D. 0-1-Tree 題解(換根dp)
阿新 • • 發佈:2021-10-24
題目連結
題目思路
好久沒寫換根dp了。。
設\(dp[i][0]\)表示以\(1\)為根節點,以\(i\)為子樹出現的\(0000/0000111\)這樣的情況
設\(dp[i][1]\)表示以\(1\)為根節點,以\(i\)為子樹出現的\(11111\)這樣的情況
然後再換根計算延申到父親的節點即可
\(dp\)方程有點複雜,需要仔細思考
程式碼
不擺爛了,寫題#include<bits/stdc++.h> #define fi first #define se second #define debug cout<<"I AM HERE"<<endl; using namespace std; typedef long long ll; const int maxn=2e5+5,inf=0x3f3f3f3f,mod=1e9+7; const double eps=1e-6; int n; ll ans[maxn]; int head[maxn],cnt; int dp[maxn][2]; int sum[maxn][2]; struct edge{ int to,next,w; }e[maxn<<2]; void add(int u,int v,int w){ e[++cnt]={v,head[u],w}; head[u]=cnt; } void dfs(int u,int fa){ for(int i=head[u];i;i=e[i].next){ int to=e[i].to; if(to==fa) continue; dfs(to,u); if(e[i].w==1){ dp[u][1]+=1+dp[to][1]; }else if(e[i].w==0){ dp[u][0]+=1+dp[to][0]+dp[to][1]; } } } void solve(int u,int fa){ ans[u]=sum[u][0]+dp[u][0]+dp[u][1]+sum[u][1]; for(int i=head[u];i;i=e[i].next){ int to=e[i].to; if(to==fa) continue; if(e[i].w==1){ sum[to][1]=sum[u][1]+1+dp[u][1]-dp[to][1]-1; }else if(e[i].w==0){ sum[to][0]=1+sum[u][0]+sum[u][1]+dp[u][0]-dp[to][0]-dp[to][1]-1+dp[u][1]; } solve(to,u); } } signed main(){ scanf("%d",&n); for(int i=1,u,v,w;i<=n-1;i++){ scanf("%d%d%d",&u,&v,&w); add(u,v,w); add(v,u,w); } dfs(1,1); solve(1,1); ll pr=0; for(int i=1;i<=n;i++){ pr+=ans[i]; } printf("%lld\n",pr); return 0; }