The Preliminary Contest for ICPC Asia Shenyang 2019

 D. Fish eating fruit

#include<bits/stdc++.h>
using namespace std;

#define ll long long
const int maxn=2e4+10;
const int mod=1e9+7;
int head[maxn],ver[maxn],nxt[maxn];
ll edge[maxn];
ll s1[maxn][3],s2[maxn][3];
ll size1[maxn][3],size2[maxn][3];
int tot=1;

void add(int u,int v,ll w)
{
    edge[++tot]=w;
    ver[tot]=v;
    nxt[tot]=head[u];
    head[u]=tot;
}
void dfs1(int u,int fa)
{
    size1[u][0]=1;
    for(int i=head[u]; i; i=nxt[i])
    {
        int v=ver[i];
        ll w=edge[i];
        if(v==fa) continue;
        dfs1(v,u);
        for(int j=0; j<3; j++)
        {
            s1[u][(w+j)%3]=(s1[u][(w+j)%3]+s1[v][j]+size1[v][j]*w)%mod;
            size1[u][(w+j)%3]=(size1[u][(w+j)%3]+size1[v][j])%mod;
        }
    }
}

void dfs2(int u,int fa)
{
    for(int i=head[u]; i; i=nxt[i])
    {
        int v=ver[i];
        if(v==fa) continue;
        ll w=edge[i];
        for(int j=0; j<3; j++)
        {
            size2[v][(j+w)%3]=((size1[u][j]-size1[v][((j-w)%3+3)%3]+size2[u][j])%mod+mod)%mod;
            s2[v][(j+w)%3]=(s1[u][j]-s1[v][((j-w)%3+3)%3]-w*(size1[v][((j-w)%3+3)%3])+s2[u][j]+size2[v][(j+w)%3]*w+mod)%mod;
        }
        dfs2(v,u);
    }
}
int main()
{
    int n;
    while(~scanf("%d",&n))
    {
        tot=1;
        memset(head,0,sizeof(head));
        memset(s1,0,sizeof(s1));
        memset(s2,0,sizeof(s2));
        memset(size1,0,sizeof(size1));
        memset(size2,0,sizeof(size2));
        for(int i=1; i<=n-1; i++)
        {
            int u,v;
            ll w;
            scanf("%d%d%lld",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        ll ans[3];
        ans[0]=ans[1]=ans[2]=0;
        dfs1(0,-1);
        dfs2(0,-1);
        for(int i=0; i<n; i++)
            for(int j=0; j<3; j++)
                ans[j]=(ans[j]+s1[i][j]+s2[i][j])%mod;
        printf("%lld %lld %lld
",ans[0],ans[1],ans[2]);
    }
}
原文地址:https://www.cnblogs.com/dongdong25800/p/11545067.html