[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;}