[洛谷P3806]【模板】点分治1

题目大意：给定一棵有$n$个点的树，$m$个询问树上距离为$k$的点对是否存在。

题解：离线，点分治

卡点：读入边的时候不知道我在干什么。。。

 

C++ Code：

#include <cstdio>
#include <algorithm>
#define maxn 100010
#define maxk 111
const int inf = 0x3f3f3f3f;
inline int max(int a, int b) {return a > b ? a : b;}

int head[maxn], cnt;
struct Edge {
	int to, nxt, w;
} e[maxn << 1];
inline void add(int a, int b, int c) {
	e[++cnt] = (Edge) {b, head[a], c}; head[a] = cnt;
	e[++cnt] = (Edge) {a, head[b], c}; head[b] = cnt;
}

bool vis[maxn];
namespace Center_of_Gravity {
	#define n __nodenum
	int root, MIN, n;
	int sz[maxn];
	void __getroot(int u, int fa) {
		sz[u] = 1;
		int MAX = 0;
		for (int i = head[u]; i; i = e[i].nxt) {
			int v = e[i].to;
			if (v != fa && !vis[v]) {
				__getroot(v, u);
				sz[u] += sz[v];
				MAX = max(sz[v], MAX);
			} 
		}
		MAX = max(n - sz[u], MAX);
		if (MAX < MIN) MIN = MAX, root = u;
	}
	int getroot(int rt, int nodenum) {
		n = nodenum;
		MIN = inf;
		__getroot(rt, 0);
		return root;
	}
	#undef n
}
using Center_of_Gravity::getroot;

int S[maxn], tot;
void getlist(int u, int fa, int val) {
	S[++tot] = val;
	for (int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if (v != fa && !vis[v]) getlist(v, u, val + e[i].w);
	}
}
int calc(int u, int val, int k) {
	tot = 0;
	getlist(u, 0, val);
	std::sort(S + 1, S + tot + 1);
	int l = 1, r = tot, res = 0;
	while (l < r) {
		if (S[l] + S[r] > k) r--;
		else if (S[l] + S[r] < k) l++;
		else {
			int tmpl = l, tmpr = r;
			while (S[tmpl] == S[l] && tmpl <= tot) tmpl++;
			while (S[tmpr] == S[r] && tmpr) tmpr--;
			int x = r - tmpr, y = tmpl - l;
			if (S[l] != S[r]) res += x * y;
			else res += x * (x - 1) >> 1;
			l = tmpl, r = tmpr;
		}
	}
	return res;
}
int V[maxn], K[maxk];
int n, m;
void solve(int u) {
	vis[u] = true;
	for (int i = 0; i < m; i++) V[i] += calc(u, 0, K[i]);
	for (int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if (!vis[v]) {
			for (int j = 0; j < m; j++) V[j] -= calc(v, e[i].w, K[j]);
			solve(getroot(v, Center_of_Gravity::sz[v]));
		}
	}
}

int main() {
	scanf("%d%d", &n, &m);
	for (int i = 1, a, b, c; i < n; i++) {
		scanf("%d%d%d", &a, &b, &c);
		add(a, b, c);
	}
	for (int i = 0; i < m; i++) scanf("%d", K + i);
	solve(getroot(1, n));
	for (int i = 0; i < m; i++) puts(V[i] ? "AYE" : "NAY");
	return 0;
}