[BZOJ3399]連通塊計數

題目描述

給出一棵n個點的樹，每個點有一個權值ai。從這棵樹上選出一個點集，使得選出的點連通，且滿足點集中的點的權值最大值與最小值之差不超過k，問有多少種選點集的辦法。

兩種選點集的辦法不同當且僅當點集中的點的標號不同。

輸入

第一行，包含兩個整數n,k。（0≤n,k≤2000）

第二行，包含n個整數a1, a2, · · · , an。（0≤ai≤2000）

接下來n − 1行，每行包含兩個正整數u, v，表示u, v兩點間有一條邊。

輸出

僅輸出一行，包含一個數，表示選點集的辦法。

這個數可能很大，輸出時對998244353取模。

樣例輸入

4 1

2 1 3 2

1 2

1 3

3 4

樣例輸出

8

來源 Quack

題解：

將點按權值排序，小的在前。

每次以一個點爲根DFS，找不比根小且不比根大k+1的點。這樣可避免重複。

#include<cstdio>
#include<algorithm>
using namespace std;
const int N=2005;
const int Mod=998244353;

int n, k, sum, a[N], s, e;
struct nodep{ int v, id; } p[N];

int fir[N], ecnt;
struct node{ int e, next; } edge[N<<1];
void Link( int s, int e ) {
    edge[++ecnt].e=e; edge[ecnt].next=fir[s]; fir[s]=ecnt;
    edge[++ecnt].e=s; edge[ecnt].next=fir[e]; fir[e]=ecnt;
}

bool cmp( nodep a, nodep b ) {
	if( a.v!=b.v ) return a.v<b.v;
	return a.id<b.id;
}

bool vis[N];
int DFS( int r, int fa, int v ) {
	int res=1;
	for( int i=fir[r]; i; i=edge[i].next )
		if( edge[i].e!=fa && !vis[ edge[i].e ] && a[ edge[i].e ]<=v )
			res=1ll*res*DFS( edge[i].e, r, v )%Mod;//將各子樹總點數相乘
	if( fa!=-1 ) res++;//非根點要算上自己
	return res;
}

int main() {
	scanf( "%d%d", &n, &k );
	for( int i=1; i<=n; i++ ) {
		scanf( "%d", &a[i] );
		p[i].v=a[i]; p[i].id=i;
	}
	for( int i=1; i<n; i++ ) {
		scanf( "%d%d", &s, &e );
		Link( s, e );
	}
	sort( p+1, p+n+1, cmp );
	
	for( int i=1; i<=n; i++ ) 
		if( !vis[ p[i].id ] ) {
			vis[ p[i].id ]=1;
			sum=( 1ll*sum+1ll*DFS( p[i].id, -1, p[i].v+k ) )%Mod;
		}
	printf( "%d\n", sum );
	return 0;
}