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\n",ans[0],ans[1],ans[2]); } }