【2019年沈阳网络赛D题】【Fish eating fruit】【树形dp(转移写法)】

题目链接：
https://nanti.jisuanke.com/t/41403
题意：
题目给了一个 $N$个点 $N-1$条边的树形图
现在每次出行的交通工具是鱼，而鱼对于路径长度有不同的喜好
如果总路径长度是3的倍数，那么鱼需要路径长度数量的 榴莲
如果总路径长度%3=1，那么鱼需要路径长度数量的 木瓜
如果总路径长度%3=2，那么鱼需要路径长度数量的 牛奶果
现在问，从图上的每个点 $i$到达除了 $i$之外的所有点（ $N*(N-1)$条路径），总共需要花费多少榴莲、木瓜、牛奶果
题解：
这里暴力解法是以每个点为根节点，然后算答案，复杂度为 $O(n^2)$，显然会T
那么考虑根节点的转移，也就是以点 $1$为根节点，然后从点 $1$开始转移树根(转移树根复杂度不能太高，要那种快速计算的才行)，计算答案，复杂度为 $O(n)$
如果没有见过这种类型的题目，可以看下洛谷P2986(这题会比较简单)

这里只需要两次 $dfs$即可完成计算：
第一次 $dfs$计算以点 $u$为根节点所得到的花费
第二次 $dfs$计算从 $u$转移到其他点作为根时得到的花费，只要不断地转移，直到遍历所有点（即所有点都做了一次树根），即可得到答案

这里我记 $dis[i]$为点 $i$的子节点到达点 $i$的路径之和
显然，若记子节点到达点 $i$的长度为 $d$，那么这样的路径可以分成三种
① $d\%3=0$
② $d\%3=1$
③ $d\%3=2$
不妨分别记做 $dis[i][0],dis[i][1],dis[i][2]$
然后我们再记录路径的数量 $num[i][0],num[i][1],num[i][2]$
记当前点为 $u$，子节点为 $son$，他们之间的路径长度为 $w$，那么我们就有了状态转移方程

num[son][0]++;
for (int j = 0; j <= 2; j++) {
		dis[u][(w + j) % 3] += num[son][j] * w + dis[son][j];
		num[u][(w + j) % 3] += num[son][j];
	}

举一个简单的例子，如果原来这条路属于③，再加上一条长度为 $2$的路，那么这条路将会划分到分类②当中，因为 $(2+2)\%3=1$
而为什么这里 $num[son][0]$需要加一呢？因为原来的 $num[son][j]$里面并不包含子节点 $son$到点 $i$的路，所以这里加上去
搞懂这个之后，我们的第一个dfs就能很轻松的写出来了

void dfs1(int u, int fa, int d) {
	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 = e[i].nxt)
		if (e[i].to != fa) {
			int son = e[i].to; ll w = e[i].w;
			dfs1(son, u, d + w);
			num[son][0]++;
			for (int j = 0; j <= 2; j++) {
				dis[u][(w + j) % 3] += num[son][j] * w + dis[son][j];
				num[u][(w + j) % 3] += num[son][j];
			}
		}
}

然后就是第二个dfs
本来我们是以点 $i$为根，那么如何快速计算以 $i$的子节点 $son$为根的答案呢？

假如我们这里将树根从点1转移到点2，那么其实，只是点 $1,3$多走了一段路1->2，点 $2,4,5$少走了一段路，然后考虑一下转移即可

void dfs2(int u, int fa) {
	for (int i = head[u]; i; i = e[i].nxt)
		if (e[i].to != fa) {
			int son = e[i].to; ll w = e[i].w;
			//记录下原来的dis,num  后面需要还原
			ll d1[10], d2[10], n1[10], n2[10];
			for (int j = 0; j <= 2; j++)d1[j] = dis[u][j], d2[j] = dis[son][j], n1[j] = num[u][j], n2[j] = num[son][j];

			//因为现在树根是son，原来的树根u变为son的子节点
			//所以原来的树根u的路径数量要减去son的路径数量(也就是dfs1里面倒过来处理一遍)
			for (int j = 0; j <= 2; j++) {
				dis[u][(w + j) % 3] -= num[son][j] * w + dis[son][j];
				num[u][(w + j) % 3] -= num[son][j];
			}
			num[son][0]--;

			//同dfs1里面，就是父子交换
			num[u][0]++;
			for (int j = 0; j <= 2; j++) {
				dis[son][(w + j) % 3] += num[u][j] * w + dis[u][j];
				num[son][(w + j) % 3] += num[u][j];
			}

			fuck(son);//统计答案
			dfs2(son, u);//递归遍历

			//还原
			for (int j = 0; j <= 2; j++)dis[u][j] = d1[j], dis[son][j] = d2[j], num[u][j] = n1[j], num[son][j] = n2[j];

		}
}

最后上总代码:

#include<bits/stdc++.h>
#define INF	0x3f3f3f3f
#define sz	sizeof
using namespace std;
typedef long long			ll;
typedef unsigned long long	ull;
typedef pair<int, int>		pii;
const int MAX = 1e4 + 10;
const int mod = 1e9 + 7;

struct edge {
	int nxt, to;
	ll w;
}e[MAX << 1];

int tot, head[MAX];

void add(int u, int v, ll w) {
	e[++tot].nxt = head[u];
	e[tot].to = v;
	e[tot].w = w;
	head[u] = tot;
}

int N;
ll ans[10], dis[MAX][3], num[MAX][3];

void fuck(int u) {//统计答案
	for (int i = 0; i <= 2; i++)
		(ans[i] += dis[u][i]) %= mod;
}

void dfs1(int u, int fa, int d) {
	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 = e[i].nxt)
		if (e[i].to != fa) {
			int son = e[i].to; ll w = e[i].w;
			
			dfs1(son, u, d + w);
			
			num[son][0]++;
			for (int j = 0; j <= 2; j++) {
				dis[u][(w + j) % 3] += num[son][j] * w + dis[son][j];
				num[u][(w + j) % 3] += num[son][j];
			}
		}
}

void dfs2(int u, int fa) {
	for (int i = head[u]; i; i = e[i].nxt)
		if (e[i].to != fa) {
			int son = e[i].to; ll w = e[i].w;
			//记录下原来的dis,num  后面需要还原
			ll d1[10], d2[10], n1[10], n2[10];
			for (int j = 0; j <= 2; j++)d1[j] = dis[u][j], d2[j] = dis[son][j], n1[j] = num[u][j], n2[j] = num[son][j];

			//因为现在树根是son，原来的树根u变为son的子节点
			//所以原来的树根u的路径数量要减去son的路径数量(也就是dfs1里面倒过来处理一遍)
			for (int j = 0; j <= 2; j++) {
				dis[u][(w + j) % 3] -= num[son][j] * w + dis[son][j];
				num[u][(w + j) % 3] -= num[son][j];
			}
			num[son][0]--;

			//同dfs1里面，就是父子交换
			num[u][0]++;
			for (int j = 0; j <= 2; j++) {
				dis[son][(w + j) % 3] += num[u][j] * w + dis[u][j];
				num[son][(w + j) % 3] += num[u][j];
			}

			fuck(son);//统计答案
			dfs2(son, u);//递归遍历

			//还原
			for (int j = 0; j <= 2; j++)dis[u][j] = d1[j], dis[son][j] = d2[j], num[u][j] = n1[j], num[son][j] = n2[j];

		}
}


int main() {
	while (~scanf("%d", &N)) {
		tot = 0;
		memset(ans, 0, sz(ans));
		memset(head, 0, sz(head));
		for (int i = 1; i < N; i++) {
			int u, v; ll w;
			scanf("%d%d%lld", &u, &v, &w);
			u++, v++;
			add(u, v, w); add(v, u, w);
		}
		dfs1(1, 0, 0);//随便找一个点为树根
		fuck(1);//统计答案
		dfs2(1, 0);//从刚才找的树根开始转移
		printf("%lld %lld %lld\n", ans[0], ans[1], ans[2]);
	}
	return 0;
}