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
",ans);
    return 0;
}