P3806 【模板】点分治1

#include<bits/stdc++.h>

using namespace std;
const int maxn = 20005;

int head[maxn], cnt = 1, n, m, d[maxn], siz[maxn], f[maxn], vis[maxn];
struct Node{
	int v, nxt, w;
} G[maxn];
int root, ans, can[maxn], tot, dep[maxn];
void ins(int u, int v, int w) {
	G[cnt] = (Node) {v, head[u], w}; head[u] = cnt++;
}
int q[maxn];

void get_rt(int x, int fa) {
	siz[x] = 1; f[x] = 0;
	for (int i = head[x]; i; i = G[i].nxt) {
		int v = G[i].v;
		if(vis[v] || v == fa) continue;
		get_rt(v, x); 
		siz[x] += siz[v];
		f[x] = max(f[x], siz[v]);
	} f[x] = max(f[x], ans - siz[x]);
	if(f[x] < f[root]) root = x;
}
void get_dp(int x, int fa) {
	d[++tot] = dep[x];
	for (int i = head[x]; i; i = G[i].nxt) {
		int v = G[i].v;
		if(vis[v] || v == fa) continue;
		dep[v] = dep[x] + G[i].w;
		get_dp(v, x);
	}
}
void calc(int x, int w, int sig) {
	tot = 0; dep[x] = w;
	get_dp(x, 0);
	sort(d + 1, d + tot + 1);
	for (int i = 1; i <= m; ++i) {
		int l = 1, r = tot, res1 = 0, res2 = 0;
		while(l < r) {
			if(d[l] + d[r] <= q[i]) res1 += (r - l), l++;
			else r--;
		} l = 1, r = tot;
		while(l < r) {
			if(d[l] + d[r] <= q[i] - 1) res2 += (r - l), l++;
			else r--;
		} can[i] += (res1 - res2) * sig;
	}
		
}
void work(int x) {
	calc(x, 0, 1); vis[x] = 1;
	for (int i = head[x]; i; i = G[i].nxt) {
		int v = G[i].v;
		if(vis[v]) continue;
		calc(v, G[i].w, -1);
		ans = siz[v]; root = 0; 
		get_rt(v, 0); work(root);
	}
}
int main() {
	scanf("%d%d", &n, &m);
	for (int i = 1; i <= n - 1; ++i) {
		int x, y, z; scanf("%d%d%d", &x, &y, &z);
		ins(x, y, z); ins(y, x, z);
	}
	for (int i = 1; i <= m; ++i) scanf("%d", &q[i]);
	root = 0, ans = n; f[0] = 0x7fffff;
	get_rt(1, 0); work(root);
	for (int i = 1; i <= m; ++i) {
		can[i] > 0 ?printf("AYE\n") :printf("NAY\n");
	}
	return 0;
}