5913. 【NOIP2018模擬10.19】林下風氣

Description
裡口福因有林下風氣，帶領全國各地高校掀起了一股AK風，大家都十分痴迷於AK。裡口福為了打擊大家的自信心，出了一道自以為十分困難的題目。
裡口福有一棵樹，第i個節點上有點權ai，他的問題就是這棵樹中有多少個不同的連通塊滿足連通塊的最大值與最小值之差=k，兩個連通塊不同當且僅當至少存在一個節點在一個連通塊中出現而另一個連通塊中沒有出現。
痴迷於AK的你馬上接下這道題目，在裡口福狂妄的笑聲中，你切掉這道題的決心更加堅定了，現在就差你的程式碼了。

Input
第一行兩個整數n,k，表示樹的大小以及題目中的k。
第二行n個整數，第i個整數表示ai。
接下來n-1行，每行兩個整數x,y表示樹邊(x,y)。

Output
一行一個整數，表示答案，答案對19260817取模。

Sample Input
5 3
1 2 3 4 5
1 2
1 3
2 4
2 5

Sample Output
4

Data Constraint
對於30%的資料，n<=22
對於另外20%的資料，樹是一條鏈
對於另外20%的資料，ai只有0和1兩種
對於100%的資料，N<=3333,0<=ai<=N,K>=0

分析：設f[i]表示以i為根的子樹中最大值為w[i]的最大值-最小值<=k的聯通塊數量，f[i] = ∏f[j]+1不難發現，如果<=k的答案減去<=k-1的答案即為==k的答案。

程式碼

#include <cstdio>
#define N 10000
#define ll long long
#define mo 19260817
using namespace std;

struct arr
{
	int to,nxt;
}a[N];
int n,k,val[N];
int ls[N],l;

void add(int x, int y)
{
	a[++l].to = y;
	a[l].nxt = ls[x];
	ls[x] = l;
} 

ll dfs(int x, int fa, int root, int k)
{
	ll sum = 1;
	for (int i = ls[x]; i; i = a[i].nxt)
		if (a[i].to != fa && val[root] >= val[a[i].to] && val[root] - val[a[i].to] <= k && (root < a[i].to || val[root] != val[a[i].to])) 
			sum = (sum * (dfs(a[i].to, x, root, k) + 1)) % mo;
	return sum % mo;
}

int main()
{
	freopen("lkf.in","r",stdin);
	freopen("lkf.out","w",stdout);
	scanf("%d%d", &n, &k);
	for (int i = 1; i <= n; i++) scanf("%d", &val[i]);
	for (int i = 1; i < n; i++)
	{
		int x, y;
		scanf("%d%d", &x, &y);
		add(x, y);
		add(y, x);
	}
	ll ans1 = 0;
	ll ans2 = 0;
	for (int i = 1; i <= n; i++)
	{
		ans1 = (ans1 + dfs(i, 0, i, k)) % mo;
		if (k) ans2 = (ans2 + dfs(i, 0, i, k - 1)) % mo;
	}
	printf("%lld", (ans1 + mo - ans2) % mo);
}