E. Alternating Tree 題解(樹形dp)
阿新 • • 發佈:2021-08-06
題目連結
題目思路
比較經典的求貢獻問題
求出每個點子樹內和子樹外和他相距奇數和偶數的點
子數外的有點小細節
然後分類討論討論求貢獻
程式碼
卷也卷不過,躺又躺不平#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 ll INF=0x3f3f3f3f3f3f3f3f; const int maxn=2e5+5,inf=0x3f3f3f3f,mod=1e9+7; const double eps=1e-6; int n; vector<int> g[maxn]; ll dp1[maxn][2],dp2[maxn][2]; ll sz[maxn]; ll a[maxn]; ll ans=0; void dfs1(int son,int fa){ sz[son]=1; dp1[son][0]=1;//包括自己 for(int i=0;i<g[son].size();i++){ int x=g[son][i]; if(x==fa) continue; dfs1(x,son); sz[son]+=sz[x]; dp1[son][0]+=dp1[x][1]; dp1[son][1]+=dp1[x][0]; } } void dfs2(int son,int fa){ for(int i=0;i<g[son].size();i++){ int x=g[son][i]; if(x==fa) continue; dp2[x][0]=dp2[son][1]+(dp1[son][1]-dp1[x][0]); dp2[x][1]=dp2[son][0]+(dp1[son][0]-dp1[x][1]); dfs2(x,son); } } void dfs3(int son,int fa){ ans+=a[son]*sz[son]; ans=(ans%mod+mod)%mod; for(int i=0;i<g[son].size();i++){ int x=g[son][i]; if(x==fa) continue; ans+=a[son]*dp1[x][1]%mod*(sz[son]-sz[x]); ans-=a[son]*dp1[x][0]%mod*(sz[son]-sz[x]); ans=(ans%mod+mod)%mod; dfs3(x,son); } } int main(){ cin>>n; for(int i=1;i<=n;i++){ cin>>a[i]; } for(int i=1,u,v;i<=n-1;i++){ cin>>u>>v; g[u].push_back(v); g[v].push_back(u); } dfs1(1,-1); dfs2(1,-1); dfs3(1,-1); // 子樹外一個點 子樹內一個點 for(int i=1;i<=n;i++){ ll add=0; add+=a[i]*dp1[i][0]%mod*(n-sz[i]); add-=a[i]*dp1[i][1]%mod*(n-sz[i]); add+=a[i]*dp2[i][0]%mod*sz[i]; add-=a[i]*dp2[i][1]%mod*sz[i]; ans=((ans+add)%mod+mod)%mod; } printf("%lld\n",ans); return 0; }