P4149 [IOI2011]Race

#include<bits/stdc++.h>

using namespace std;
const int maxn = 400005;

int n, k, root, tot, cnt = 1, head[maxn], f[maxn];
int siz[maxn], ans[1000005], sum, vis[maxn], deep[maxn], dis[maxn];
struct Node{
	int v, nxt, w;
} G[maxn]; 
struct DATA {
	int dep, l;
} a[maxn];
void ins(int u, int v, int w) {
	G[cnt] = (Node) {v, head[u], w}; head[u] = cnt++;
}
int read() {
	int x = 0, f = 1; char ch = getchar();
	while(ch < ‘0‘ || ch > ‘9‘) { if(ch == ‘-‘) f = -1; ch = getchar();}
	while(ch >= ‘0‘ && ch <= ‘9‘) { x = x * 10 + ch - 48; ch = getchar();}
	return x * f;
}

void Getrt(int x, int fa) {
	f[x] = 0; siz[x] = 1;
	for (int i = head[x]; i; i = G[i].nxt) {
		int v = G[i].v;
		if(v == fa || vis[v]) continue;
		Getrt(v, x);
		siz[x] += siz[v];
		f[x] = max(f[x], siz[v]);
	} f[x] = max(f[x], sum - siz[x]);
	if(f[x] < f[root]) root = x;
}
bool cmp(DATA a, DATA b) {
	return a.l < b.l;
}
void Getdp(int x, int fa) {
	a[++tot] = (DATA) {deep[x], dis[x]};
	for (int i = head[x]; i; i = G[i].nxt) {
		int v = G[i].v;
		if(vis[v] || v == fa) continue;
		deep[v] = deep[x] + 1;
		dis[v] = dis[x] + G[i].w;
		Getdp(v, x);
	}
}
void calc(int x, int w, int ww, int sig) {
	deep[x] = ww; dis[x] = w; tot = 0; Getdp(x, 0);
	sort(a + 1, a + tot + 1, cmp);
	int l = 1, r = tot;
	for (; l <= r; ++l) {
		while(l < r && a[l].l + a[r].l > k) --r;
		for (int i = r; a[i].l + a[l].l == k; --i) ans[a[i].dep + a[l].dep] += sig; 
	}
}
void work(int x) {
	vis[x] = 1; calc(x, 0, 0, 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, -1);
		root = 0; sum = siz[v];
		Getrt(v, 0); work(root);
	}
}

int main() {
	n = read(); k = read();
	for (int i = 1; i <= n - 1; ++i) {
		int x, y, z; x = read(); y = read(); z = read(); 
		x++; y++;
		ins(x, y, z); ins(y, x, z);
	}
	root = 0; f[0] = 0x7fffff; sum = n;
	Getrt(1, 0); work(root);
	for (int i = 1; i <= n; ++i) {
		if(ans[i]) {
			printf("%d\n", i); return 0;
		}
	} printf("-1");
	return 0;
}