【树形dp】连通块计数

题目描述

给出一棵n个点的树，每个点有一个权值a。从这棵树上选出一个点集，使得选出的点连通，且满足点集中的点的权值最大值与最小值之差不超过k，问有多少种选点集的办法。

两种选点集的办法不同当且仅当点集中的点的标号不同。

输入

第一行，包含两个整数n,k。

第二行，包含n个整数a 1, a 2, · · · , a n。

接下来n − 1行，每行包含两个正整数u, v，表示u, v两点间有一条边。

输出

仅输出一行，包含一个数，表示选点集的办法。

这个数可能很大，输出时对998244353取模。

样例输入

Copy (如果复制到控制台无换行，可以先粘贴到文本编辑器，再复制)

4 1
2 1 3 2
1 2
1 3
3 4

样例输出

8

提示

对于100%的数据，0 ≤ n, k, a i  ≤ 2000。

题解

我们可以以每一个节点为树根，然后对整棵树进行计数。

设f[i]表示以i为根的子树中符合题意的联通块的个数，则 f[i]= ∏(j ∈ s(i)) f[j]

其中s(i)表示i的儿子的集合

-----------------------------------------------------------------------------------------------------------------------------------------------------

不难发现，如果直接这样做，将会有不少联通块被重复计算。那么，怎么去重呢？

我们可以重新规定f[i]为以w[i]为某连通块中最大的权值的连通块个数，则f[i]=∏(j ∈ s(i)) f[j]，j满足：

1、w[root[i]] >= w[j]

2、w[root[i]]-w[j] <= k

感觉像是搞定了，但实际上还是有一个问题：如果权值有相同的，这种方法还是会出现重复计算的问题。

针对这个问题，我们还可以按照某一个方向计算的节点的编号。如按编号从小到大计算。这样就可以保证不会出现重复的了。

代码

#include<cstdio>
#include<vector>
using namespace std;
const int mn = 2005, mod = 998244353;
int w[mn], k, g[mn][mn], p[mn];
int dp(int u, int f, int x)
{
	int ret = 1;
	for(int i = 1, v; v = g[u][i]; ++i)
		if(v != f && w[x] >= w[v] && w[x] - w[v] <= k && (x < v || w[x] != w[v]))
			ret = 1ll * ret * (dp(v, u, x) + 1) % mod;
	return ret;
}
int main()
{
	int n, a, b, ans = 0;
	int i;
	scanf("%d%d", &n, &k);
	for(i = 1; i <= n; i++)
		scanf("%d", &w[i]);
	for(i = 1; i < n; i++)
	{
		scanf("%d%d", &a, &b);
		g[a][++p[a]] = b, g[b][++p[b]] = a;
	}
	for(i = 1; i <= n; i++)
		ans = (ans + dp(i, 0, i)) % mod;
	printf("%d", ans);
}