The Preliminary Contest for ICPC Asia Shenyang 2019 D. Fish eating fruit 树形DP 换根与二次扫描

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

代码抄借鉴于：https://blog.csdn.net/weixin_44282912/article/details/100833858

#include<iostream>
#include<vector>
#include<algorithm>
#include<string>
#include<stack>
#include<cstring>
#include<set>
#include<iterator>
#include<list>
#include<deque>
#include<queue>
#include<map>
#include<cmath>
#include<sstream>
#include<cstdio>
#include<ctime>
#include<iomanip>
//#include<unordered_set>
using namespace std;
#define lowbit(x) x&(-x)
typedef long long ll;
typedef pair<int, int> P;
const int N = 1e4 + 10;
const int M = 1e4 + 10;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 1e9 + 7;
ll dis[N][3], num[N][3];
ll ans[3];
int tot;
int head[N], ver[2 * N], edge[2 * N], Next[2 * N];
inline void addedge(int x, int y, int z)
{
	ver[++tot] = y, edge[tot] = z, Next[tot] = head[x], head[x] = tot;
	ver[++tot] = x, edge[tot] = z, Next[tot] = head[y], head[y] = tot;
}
void solve(int u)
{
	for (int i = 0; i < 3; i++)
	{
		ans[i] = (ans[i] + dis[u][i]) % mod;
	}
}
void dfs(int u, int fa)
{
	dis[u][0] = dis[u][1] = dis[u][2] = 0;
	num[u][0] = num[u][1] = num[u][2] = 0;
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i], w = edge[i];
		if (v == fa)
			continue;
		dfs(v, u);
		num[v][0]++; 
		for (int j = 0; j < 3; j++)
		{
			dis[u][(j + w) % 3] += dis[v][j] + w * num[v][j];
			num[u][(j + w) % 3] += num[v][j];
		}
	}
}
void dfs2(int u, int fa)
{
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i], w = edge[i];
		if (v == fa)
			continue;
		ll tdis[2][3], tnum[2][3];
		for (int j = 0; j < 2; j++)
		{
			for (int k = 0; k < 3; k++)
			{
				if (!j)
				{
					tdis[j][k] = dis[u][k];
					tnum[j][k] = num[u][k];
				}
				else
				{
					tdis[j][k] = dis[v][k];
					tnum[j][k] = num[v][k];
				}
			}
		}
		for (int j = 0; j < 3; j++)
		{
			dis[u][(w + j) % 3] -= num[v][j] * w + dis[v][j];
			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] += num[u][j] * w + dis[u][j];
			num[v][(w + j) % 3] += num[u][j];
		}
		solve(v);
		dfs2(v, u);
		for (int j = 0; j < 2; j++)
		{
			for (int k = 0; k < 3; k++)
			{
				if (!j)
				{
					dis[u][k] = tdis[j][k];
					num[u][k] = tnum[j][k];
				}
				else
				{
					dis[v][k] = tdis[j][k];
					num[v][k] = tnum[j][k];
				}
			}
		}
	}
}
int main()
{
	int n;
	while (scanf("%d", &n) != EOF)
	{
		tot = 0;
		memset(ans, 0, sizeof(ans));
		memset(head, 0, sizeof(head));
		for (int i = 1; i < n; i++)
		{
			int x, y, z;
			scanf("%d%d%d", &x, &y, &z);
			x++, y++;
			addedge(x, y, z);
		}
		dfs(1, 0);
		solve(1);
		dfs2(1, 0);
		printf("%lld %lld %lld\n", ans[0], ans[1], ans[2]);
	}
	return 0;
}