CodeForces - 486D Valid Sets

Solution

枚举根为 $\max$，这样既好判断 $\max -\min <=d$ 而且能保证同一种联通子图不被计算多次（另外要特判与根的权值相等的情况，只能放一个来 $\mathtt{D}\mathtt{P}$）。

然后考虑选不选子树大力 $\mathtt{D}\mathtt{P}$ 就行了。

Code

#include <cstdio>

const int mod = 1e9 + 7, N = 2010;

int n, d, w[N], head[N], nxt[N << 1], to[N << 1], cnt, ans, f[N];

int read() {
	int x = 0, f = 1; char s;
	while((s = getchar()) > '9' || s < '0') if(s == '-') f = -1;
	while(s >= '0' && s <= '9') x = (x << 1) + (x << 3) + (s ^ 48), s = getchar();
	return x * f;
}

void addEdge(const int u, const int v) {
	nxt[++ cnt] = head[u], to[cnt] = v, head[u] = cnt;
}

void DP(const int u, const int fa, const int root) {
	f[u] = 1;
	for(int i = head[u]; i; i = nxt[i]) {
		int v = to[i];
		if(v == fa) continue;
		if((w[root] > w[v] || (w[root] == w[v] && root > v)) && w[root] - w[v] <= d)
			DP(v, u, root), f[u] = (f[u] + 1ll * f[u] * f[v]) % mod;
	}
}

int main() {
	int u, v;
	d = read(), n = read();
	for(int i = 1; i <= n; ++ i) w[i] = read();
	for(int i = 1; i < n; ++ i) {
		u = read(), v = read();
		addEdge(u, v); addEdge(v, u);
	}
	for(int i = 1; i <= n; ++ i) DP(i, 0, i), ans = (ans + f[i]) % mod;
	printf("%d\n", ans);
	return 0;
}