2019icpc沈阳网络赛D Fish eating fruit

https://nanti.jisuanke.com/t/41403
推荐这篇blog，讲的极好：https://blog.csdn.net/weixin_44282912/article/details/100833858

#include<bits/stdc++.h>
using namespace std;
const int maxn=10000+100;
const int mod=1000000007;
typedef long long ll;

ll n,dis[maxn][3],num[maxn][3],ans[3];
struct Edge{
	int to,dist;
};
vector<Edge> edges;
vector<int> G[maxn];

void AddEdge(int a,int b,int c)
{
	edges.push_back((Edge){b,c});
	G[a].push_back(edges.size()-1);
	edges.push_back((Edge){a,c});
	G[b].push_back(edges.size()-1);
}

void dfs1(int u,int fa)
{	
	for(int i=0;i<G[u].size();i++)
	{
		int v=edges[G[u][i]].to,w=edges[G[u][i]].dist;
		if(v==fa)continue;
		dfs1(v,u);
		num[v][0]++;
		for(int j=0;j<3;j++)dis[u][(w+j)%3]=(dis[u][(w+j)%3]+dis[v][j]+w*num[v][j])%mod,num[u][(w+j)%3]+=num[v][j];
	}
}

void dfs2(int u,int fa)
{
	for(int j=0;j<3;j++)ans[j]=(ans[j]+dis[u][j])%mod;
	for(int i=0;i<G[u].size();i++)
	{
		int v=edges[G[u][i]].to,w=edges[G[u][i]].dist;
		if(v==fa)continue;
		for(int j=0;j<3;j++)dis[u][(w+j)%3]=(dis[u][(w+j)%3]-dis[v][j]-w*num[v][j]%mod+2*mod)%mod,num[u][(w+j)%3]-=num[v][j];
		num[v][0]--;
		num[u][0]++;
		for(int j=0;j<3;j++)dis[v][(w+j)%3]=(dis[v][(w+j)%3]+dis[u][j]+w*num[u][j])%mod,num[v][(w+j)%3]+=num[u][j];
		dfs2(v,u);
		for(int j=0;j<3;j++)dis[v][(w+j)%3]=(dis[v][(w+j)%3]-dis[u][j]-w*num[u][j]%mod+2*mod)%mod,num[v][(w+j)%3]-=num[u][j];
		num[u][0]--;
		num[v][0]++;
		for(int j=0;j<3;j++)dis[u][(w+j)%3]=(dis[u][(w+j)%3]+dis[v][j]+w*num[v][j])%mod,num[u][(w+j)%3]+=num[v][j];
	}
}

int main()
{
	//freopen("input.in","r",stdin);
	while(cin>>n)
	{
		memset(num,0,sizeof(num));
		memset(dis,0,sizeof(dis));
		memset(ans,0,sizeof(ans));
		for(int i=0;i<n;i++)G[i].clear();
		edges.clear();
		int a,b,c;
		for(int i=1;i<n;i++)scanf("%d%d%d",&a,&b,&c),AddEdge(a,b,c);
		dfs1(0,-1);
		dfs2(0,-1);
		printf("%lld %lld %lld\n",ans[0],ans[1],ans[2]);
	}
}