E. Alternating Tree 题解(树形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 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;
}
不摆烂了,写题