Fish eating fruit | 易学教程

测试地址
题意简述：
树上任意两点之间的路径按照模 3 为 012 分类，将两点间距离加和，乘 2 即为答案。

解题思路：
可以采用树上dp解决，也可以点分治，这里先给出一种树上dp做法：
dp[i][k] 表示距 i 模 3 为 k 的节点距离和。
tc[i][k] 表示距 i 模 3 为 k 的节点数目。
ans[k] 表示所有路径中模 3 为 k 的路径的总长度。
目标答案是ans[k] 。
初始状态 tc[i][0] = 1 。

如果每次只考虑所有经过根 x 的路径，并且路径上的一个端点在x的一个子树上，另一个端点在另一个子树上（其他所有情况都可以在x的祖先或者子节点被考虑到，所以这样可以包含所有情况）。
假设当前枚举到x的子节点y，之前遍历的子节点已经使得dp和tc数组更新完成，那么我们要计算的路径起点在y，终点在之前遍历过的所有子节点中。
分类讨论答案贡献：

边 x-y 对答案的贡献：设 j，k属于{0,1,2}，x 到 y 的边权为 z ，那么z对答案的贡献为tc[x][j] * tc[y][k] * z。
终点是 y 的所有路径长度的贡献：dp[y][k] * tc[x][j] * z。
起点是 x 的所有路径长度的贡献：dp[x][j] * tc[y][k] * z。

于是状态转移方程：
dp[x][(j+z)%3] += dp[y][j] + z * tc[y][j]
tc[x][(j+z)%3] += tc[y][j]

当然，在状态转移前也要更新ans数组。
代码示例：

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

const int P = 1e9+7;
const int N = 1e5;
typedef long long ll;

int head[N],ver[N],edge[N],nex[N];
int tot = 1;
void addEdge(int x,int y,int z){
	ver[++tot] = y; edge[tot] = z;
	nex[tot] = head[x]; head[x] = tot;
}

int n;
ll tc[N][4], dp[N][4], ans[4];

void dfs(int x,int fa){
	/* 利用dfs进行状态转移,x为当前子树根节点 */ 
	for(int i = head[x];i != -1;i = nex[i]){
		int y = ver[i], z = edge[i];
		if(y == fa) continue;	//y是父节点则跳过 
		dfs(y,x);					
		/* 这里统计答案 */ 
		for(int j = 0;j < 3;j++)
			for(int k = 0;k < 3;k++){
				ans[(j+k+z)%3] += (dp[x][j]*tc[y][k]%P+dp[y][k]*tc[x][j]%P)%P;
				ans[(j+k+z)%3] += z*tc[x][j]%P*tc[y][k]%P;
				ans[(j+k+z)%3] %= P;
			}
		/* 在这里转移状态 */ 
		for(int j = 0;j < 3;j++){
			dp[x][(j+z)%3] = (dp[x][(j+z)%3] + dp[y][j] + z*tc[y][j])%P;
			tc[x][(j+z)%3] = (tc[x][(j+z)%3] + tc[y][j])%P;
		}
	}
}
void solve(){
	memset(dp,0,sizeof dp);
	memset(tc,0,sizeof tc);
	memset(ans,0,sizeof ans);
	for(int i = 0;i <= n;i++) tc[i][0] = 1;
	/* 统计答案并输出 */ 
	dfs(1,0);
	for(int i = 0;i < 2;i++) printf("%lld ",ans[i]*2%P);
	printf("%lld\n",ans[2]*2%P);
}
int main(){
	while(~scanf("%d",&n)){
		memset(head,-1,sizeof head); tot = 1;
		for(int i = 1,x,y,z;i < n;i++){
			scanf("%d%d%d",&x,&y,&z);
			addEdge(x+1,y+1,z); addEdge(y+1,x+1,z);
		}
		solve();
	}
	return 0;
}

来源：https://blog.csdn.net/weixin_41162823/article/details/101264046

标签