点分治练习

P3806 【模板】点分治 1

#include <bits/stdc++.h>
using namespace std;


inline long long read() {
	char ch = getchar();
	long long f = 1,x = 0;
	while (ch > '9' || ch < '0') { if (ch == '-')f = -1; ch = getchar(); }
	while (ch >= '0' && ch <= '9') { x = (x << 1) + (x << 3) + ch - '0'; ch = getchar(); }
	return x * f;
}
//void ReadFile() {
//	FILE* stream1;
//	freopen_s(&stream1,"in.txt", "r", stdin);
//	freopen_s(&stream1,"out.txt", "w", stdout);
//}

static auto speedup = []() {ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr); return nullptr; }();

const int maxn = 1e4 + 7,INF = 1e7 + 7;
bitset<INF> judge;

int n,m,vis[maxn],sum = 0,subsum = 0,root = 0, ask[maxn],ans[maxn],siz[maxn],d[maxn],dis[maxn],cnt;
int q[INF];

struct Node {
	Node() :y(0), dis(0){}
	Node(int a,int b) :y(a), dis(b) {}
	int y,dis;
};
vector<Node> e[maxn];


void getroot(int x, int fa) {
	siz[x] = 1;
	int s = 0;
	for (auto& nx : e[x]) {
		if (nx.y == fa || vis[nx.y])continue;
		getroot(nx.y, x);
		siz[x] += siz[nx.y];
		s = max(s, siz[nx.y]);
	}
	s = max(s, sum - siz[x]);
	if (s < subsum) {
		root = x;
		subsum = s;
	}

}
void getdis(int x, int fa) {
	dis[++cnt] = d[x];
	for (auto& nx : e[x]) {
		if (nx.y == fa || vis[nx.y])continue;
		d[nx.y] = d[x] + nx.dis;
		getdis(nx.y, x);
	}
}

void calc(int x) {
	judge[0] = 1;
	int p = 0;
	for (auto& nx : e[x]) {
		if (vis[nx.y])continue;
		cnt = 0;
		d[nx.y] = nx.dis;
		getdis(nx.y, x);
		for (int j = 1; j <= cnt; j++) {
			for (int k = 1; k <= m; k++) {
				if (ask[k] >= dis[j]) {
					ans[k] |= judge[ask[k] - dis[j]];
				}
			}
		}
		for (int j = 1; j <= cnt; j++) {
			if (dis[j] < INF) {
				q[++p] = dis[j];
				judge[dis[j]] = 1;
			}
		}
	}
	for (int i = 1; i <= p; i++)judge[q[i]] = 0;
}

void divide(int x) {
	calc(x);
	vis[x] = 1;
	for (auto& nx : e[x]) {
		if (vis[nx.y])continue;
		sum = subsum = siz[nx.y];
		getroot(nx.y,0);
		divide(root);
	}
}
void slove() {
	cin >> n >> m;
	int a, b, c;
	for (int i = 1; i < n; i++) {
		cin >> a >> b >> c;
		e[a].push_back(Node{ b,c });
		e[b].push_back(Node{ a,c });
	}

	for (int i = 1; i <= m; i++) {
		cin >> ask[i];
	}
	sum = subsum = n;
	getroot(1, 0);
	getroot(root, 0);
	divide(root);
	for (int i = 1; i <= m; i++) {
		cout << (ans[i] ? "AYE" : "NAY") << endl;
	}
}
int main() {
	slove();
	return 0;
}