題目描述

給出一棵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);
}