D. Fish eating fruit

树形dp

题:https://nanti.jisuanke.com/t/41403

题意：求任意俩点之间距离之和模3后的三个结果的总数（原距离之和）

#include<bits/stdc++.h>
using namespace std;
#define pb push_back
typedef long long ll;
const int M=1e4+4;
const int mod=1e9+7;
struct node{
    int v;
    ll w;
};
ll C[M][3],S[M][3],ans[M];
//C[i][j]:表示以i为根，然后路径消耗取模后为j的路径数
//S[i][j]:表示以i为根，路径消耗取模后为j的路径总消耗 
vector<node>e[M];;
void dfs(int u,int f,ll pre){
    C[u][0]=C[u][1]=C[u][2]=0;
    S[u][0]=S[u][1]=S[u][2]=0;
    int len=e[u].size();
    for(int i=0;i<e[u].size();i++){
        int v=e[u][i].v;
        if(v==f)
            continue;
        dfs(v,u,e[u][i].w);
        //算跨越跟节点的贡献
        for(int p=0;p<3;p++){
            for(int j=0;j<3;j++)
                for(int k=0;k<3;k++)
                    if(p==(j+k)%3) 
                        ans[p]=(ans[p]+S[u][j]*C[v][k]%mod+C[u][j]*S[v][k]%mod)%mod;
        }
        for(int j=0;j<3;j++){
            C[u][j]=(C[u][j]+C[v][j])%mod;
            S[u][j]=(S[u][j]+S[v][j])%mod;
        }
        
    }
    
    for(int i=0;i<3;i++)//算以u为跟对答案的贡献，就直接算u的每一个子树的贡献 

        ans[i]=(ans[i]+S[u][i])%mod;
    ll c[3],s[3];
    memset(c,0ll,sizeof(c));
    memset(s,0ll,sizeof(s));
    for(int i=0;i<3;i++){
        int t=(i-pre%3+3)%3;
        c[i]=(c[i]+C[u][t])%mod;
        s[i]=(s[i]+(S[u][t]+C[u][t]*pre%mod)%mod)%mod;
    }
    if(f!=0)
        c[pre%3]=(c[pre%3]+1ll)%mod,s[pre%3]=(s[pre%3]+pre)%mod;
    for(int i=0;i<3;i++)
        C[u][i]=c[i],S[u][i]=s[i];
}
int main(){
    int n;
    while(~scanf("%d",&n)){
        for(int i=0;i<=n;i++)
            e[i].clear();
        memset(S,0,sizeof(S));
        memset(C,0,sizeof(C)); 
        for(int i=1;i<n;i++){
            int u,v;
            ll w;
            for(int i=0;i<3;i++)
                ans[i]=0ll;
            scanf("%d%d%lld",&u,&v,&w);
            u++,v++;
            e[u].pb(node{v,w});
            e[v].pb(node{u,w});
        }
        dfs(1,0,0);
        printf("%lld %lld %lld\n",ans[0]*2ll%mod,ans[1]*2ll%mod,ans[2]*2ll%mod);
        
    }
    return 0;;
}

View Code