JZOJ5913. 【NOIP2018模拟10.19】林下风气

Description

里口福因有林下风气，带领全国各地高校掀起了一股AK风，大家都十分痴迷于AK。里口福为了打击大家的自信心，出了一道自以为十分困难的题目。
里口福有一棵树，第i个节点上有点权ai，他的问题就是这棵树中有多少个不同的连通块满足连通块的最大值与最小值之差=k，两个连通块不同当且仅当至少存在一个节点在一个连通块中出现而另一个连通块中没有出现。
痴迷于AK的你马上接下这道题目，在里口福狂妄的笑声中，你切掉这道题的决心更加坚定了，现在就差你的代码了。

Data Constraint

对于30%的数据，n<=22
对于另外20%的数据，树是一条链
对于另外20%的数据，ai只有0和1两种
对于100%的数据，N<=3333,0<=ai<=N,K>=0

题解

对于每一个对应的最小值，必然有一个最大值。
可以枚举一个最小值，
设$f_{i,0/1,0/1}$表示以i为节点的子树里面，i必选，最小值有没有出现过，最大值有没有出现过的方案数。
转移就很简单了，直接暴力算。

code

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <cstring>
#include <string.h>
#include <cmath>
#include <math.h>
#define N 3338
#define ll long long
#define P putchar
#define G getchar
using namespace std;
char ch;
void read(int &n)
{
	n=0;
	ch=G();
	while((ch<'0' || ch>'9') && ch!='-')ch=G();
	int w=1;
	if(ch=='-')w=-1,ch=G();
	while('0'<=ch && ch<='9')n=(n<<3)+(n<<1)+ch-'0',ch=G();
	n*=w;
}

const int mo=19260817;
int nxt[N*2],to[N*2],lst[N],tot,x,y;
int n,m,k,mx,mi,v[N],a[N],dfn[N],id,ans,f[N][2][2];

void add(int& x,int y){x=(x+y<mo?x+y:x+y-mo);}

void ins(int x,int y)
{
	nxt[++tot]=lst[x];
	to[tot]=y;
	lst[x]=tot;
}

void dfs(int x,int fa)
{
	dfn[++id]=x;
	for(int i=lst[x];i;i=nxt[i])
		if(to[i]^fa)dfs(to[i],x);
}

int main()
{
	freopen("lkf.in","r",stdin);
	freopen("lkf.out","w",stdout);
	
	read(n);read(k);
	for(int i=1;i<=n;i++)
		read(a[i]),v[i]=a[i];
	sort(v+1,v+1+n);
	m=unique(v+1,v+1+n)-v-1;
	for(int i=1;i<n;i++)read(x),read(y),ins(x,y),ins(y,x);
	dfs(1,0);
	for(int i=1;i<=m;i++)
	{
		memset(f,0,sizeof(f));
		mi=v[i];mx=mi+k;
		for(int j=n;j;j--)
		{
			x=dfn[j];
			if(a[x]==mi && a[x]!=mx)f[x][1][0]=1;
			if(a[x]==mx && a[x]!=mi)f[x][0][1]=1;
			if(a[x]>mi && a[x]<mx)f[x][0][0]=1;
			if(a[x]==mi && a[x]==mx)f[x][1][1]=1;
			for(int w=lst[x];w;w=nxt[w])
			{
				y=to[w];
				f[x][1][1]=((ll)f[x][1][1]*(f[y][0][0]+f[y][0][1]+f[y][1][0]+f[y][1][1]+1)
					+(ll)f[y][1][1]*(f[x][0][0]+f[x][0][1]+f[x][1][0])
					+(ll)f[x][0][1]*f[y][1][0]+(ll)f[x][1][0]*f[y][0][1])%mo;
				f[x][0][1]=((ll)f[x][0][1]*(f[y][0][0]+1)+(ll)f[y][0][1]*(f[x][0][0]+f[x][0][1]))%mo;
				f[x][1][0]=((ll)f[x][1][0]*(f[y][0][0]+1)+(ll)f[y][1][0]*(f[x][0][0]+f[x][1][0]))%mo;
				f[x][0][0]=(ll)f[x][0][0]*(f[y][0][0]+1)%mo;
			} 
			add(ans,f[x][1][1]);
		}
	}
	printf("%d\n",ans);
	
	return 0;
}