SPOJ 1825 Free tour II （树的点分治）

题目链接 Free tour II

题意：有$N$个顶点的树，节点间有权值， 节点分为黑点和白点。 找一条最长路径使得 路径上黑点数量不超过K个

这是树的点分治比较基本的题，涉及树上启发式合并……仰望了黄学长的博客之后稍微有点明白了（还没有很深入地理解）

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b)	for (int i(a); i <= (b); ++i)
#define dec(i, a, b)	for (int i(a); i >= (b); --i)

const int N = 200010;

int ans, n, k, m, cnt, root, sum, deep_mx, x;
int sz[N], f[N], deep[N], dis[N], tmp[N], mx[N];
bool color[N], vis[N];
vector <pair <int, int > > v[N];
vector <pair <int, int > > st;

void getroot(int x, int fa){
	sz[x] = 1; f[x] = 0;
	for (auto u : v[x]){
		int to = u.first;
		if (vis[to] || to == fa) continue;
		getroot(to, x);
		f[x] = max(f[x], sz[to]);
		sz[x] += sz[to];
	}

	f[x] = max(f[x], sum - sz[x]);
	if (f[x] < f[root]) root = x;
}

void getdis(int x, int fa){
	deep_mx = max(deep_mx, deep[x]);
	for (auto u : v[x]){
		int to = u.first;
		if (vis[to] || to == fa) continue;
		deep[to] = deep[x] + color[to];
		dis[to] = dis[x] + u.second;
		getdis(to, x);
	}
}

void getmx(int x, int fa){
	tmp[deep[x]] = max(tmp[deep[x]], dis[x]);
	for (auto u : v[x]){
		int to = u.first;
		if (vis[to] || to == fa) continue;
		getmx(to, x);
	}
}

void solve(int x){
	vis[x] = 1;
	st.clear();
	if (color[x]) --k;
	for (auto u : v[x]){
		int to = u.first;
		if (vis[to]) continue;
		deep_mx = 0;
		deep[to] = color[to];
		dis[to] = u.second;
		getdis(to, x);
		st.push_back({deep_mx, to});
	}

	sort(st.begin(), st.end());
	for (int i = 0; i < (int)st.size(); ++i){
		getmx(st[i].second, x);
		int now = 0;
		if (i != 0)
			dec(j, st[i].first, 0){
				while (now + j < k && now < st[i - 1].first)
				       ++now, mx[now] = max(mx[now], mx[now - 1]);
				if (now + j <= k) ans = max(ans, mx[now] + tmp[j]);
			}
		if (i != (int)st.size() - 1)
			rep(j, 0, (int)st[i].first)
				mx[j] = max(mx[j], tmp[j]), tmp[j] = 0;
		else
			rep(j, 0, (int)st[i].first){
				if (j <= k) ans = max(ans, max(tmp[j], mx[j]));
				tmp[j] = mx[j] = 0;
			}
	}

	if (color[x]) ++k;
	for (auto u : v[x]){
		int to = u.first;
		if (vis[to]) continue;
		root = 0;
		sum = sz[to];
		getroot(to, x);
		solve(root);
	}
}

int main(){

	scanf("%d%d%d", &n, &k, &m);
	rep(i, 1, m){
		scanf("%d", &x);
		color[x] = 1;
	}

	rep(i, 1, n - 1){
		int x, y, z;
		scanf("%d%d%d", &x, &y, &z);
		v[x].push_back({y, z});
		v[y].push_back({x, z});
	}

	sum = n; f[0] = n;
	getroot(1, 0);
	solve(root);
	printf("%d\n", ans);
	return 0;
}