点分治入门

一棵n个顶点的树, 每条边有一个长度（[1,1000]中的正整数）, 定义dis(u, v)是u和v之间的树上距离. 给定
整数k, 每个对(u, v)称之为合法如果dis(u, v)<=k,写一个程序计算树上有多少合法对?
  
【输入】
多样例. 每个样例开始于n，k， n<=1w, 然后n-1行每行包含u,v,l 表示u和v之间有边长为l
输入以0,0 结束
  
【输出】
答案
  
【样例输入】
5 4  
1 2 3  
1 3 1  
1 4 2  
3 5 1  
0 0  
  
【样例输出】
8  
  
【限制】
1s

//#include "stdafx.h"
#pragma comment(linker, "/STACK:1024000000,1024000000")
#include <stdio.h>
#include <algorithm>
#include <string.h>
using namespace std;
//#define LOCAL
const int maxn = 1e4+5, inf = 0x3f3f3f3f;
int n,k, head[maxn], cnt, dis[maxn], zx, mb, ans, tot, sz[maxn]; // zx是重心
bool v[maxn]; // v[i]=j表明i节点已经作为树根了
struct Arc  
{  
	int from, to, nxt, len;
} g[maxn << 1];
  
void addarc(int from, int to, int len)  
{
	g[cnt].from = from, g[cnt].to = to, g[cnt].len = len, g[cnt].nxt = head[from];
	head[from] = cnt++;
}
  
void mmax(int &a, int b)  
{
	if (a < b)
	{
		a = b;
	}
}
  
void dfs2(int cur, int fa) // 得到各个节点的子树大小
{
	sz[cur] = 1;
	for (int h = head[cur], to; ~h; h = g[h].nxt)
	{
		to = g[h].to;
		if (to ^ fa && !v[to])
		{
			dfs2(to, cur);
			sz[cur] += sz[to];
		}
	}
}
  
void dfs1(int cur, int fa, int nn) // 求树的重心, 返回cur子树大小. nn是求重心的树的总的节点个数
{
	int ret = 1, b = 0;
	for (int h = head[cur], to; ~h; h = g[h].nxt)
	{
		to = g[h].to;
		if (to ^ fa && !v[to])
		{
			dfs1(to, cur, nn);
			ret += sz[to];
			mmax(b, sz[to]);
		}
	}
	mmax(b, nn - ret);
	if (b < mb)
	{
		mb = b;
		zx = cur;
	}
}
  
void dfs3(int cur, int d, int fa)  
{
	dis[tot++] = d;
	for (int h = head[cur], to, len; ~h; h = g[h].nxt)
	{
		to = g[h].to;
		len = g[h].len;
		if (to ^ fa && !v[to])
		{
			dfs3(to, d + len, cur);
		}
	}
}
  
int kk(int cur, int d, int fa) // 返回cur子树中所有(a, b)的点对数, 其中(a,b)点对满足 (dis[a]+d)+(dis[b]+d)<=k, 注意,(a,b)和(b,a)视作相同点对哈，不能重复计数~ 其中dis[a]是a到cur的距离
{
	tot = 0;
	dfs3(cur, d, fa); // 获取cur子树中所有节点a的dis[a]+d的值,存入dis[0, tot)数组中
	sort(dis, dis + tot);
	int ret = 0, i = 0, j = tot - 1; // 求dis[0,..,tot) 中满足dis[a]+dis[b] <=k的(a,b)对的数量，i,j是双指针
	while(i < j)
	{
		if (dis[i] + dis[j] <= k)
		{
			ret += j - i; // 则 dis[i]和dis[i+1,...,j]都满足
			++i;
		}
		else  
		{
			--j;
		}
	}
	return ret;
}
  
void dfs(int cur, int fa) // 点分治
{
	mb = inf; // 用于求重心的
	dfs2(cur, fa); // 确定cur子树的大小sz[cur], 注意, 因为重心（即根节点）在不断变化中, 所以每次都要重新处理出sz来
	dfs1(cur, fa, sz[cur]); // 求重心zx, 它将是新的树根
	v[zx] = true;
	int t = zx; // zx是全局变量, 随着递归会变, 故而用t缓存,下面都用t
	ans += kk(t, 0, 0); // 这里传入的fa是0, 并不能阻止朝着zx去递归, 所以必须在dfs3方法中加入 !v[to]的判断(见72行)
	for (int h = head[t], to, len; ~h; h = g[h].nxt) // 注意, 树根从cur变成了t=zx
	{
		to = g[h].to;
		len = g[h].len;
		if (!v[to]) // 防止回溯到t=zx去(因为v[zx]刚刚被赋予true)
		{
			ans -= kk(to, len, t);
			dfs(to, t);
		}
	}
}
  
signed main()  
{
#ifdef LOCAL
	freopen("d:\\data.in","r",stdin);
//	freopen("d:\\my.out", "w", stdout);
#endif  
	while(scanf("%d%d", &n, &k), n)
	{
		memset(head, -1, sizeof(head));
		memset(v, 0, sizeof(v));
		cnt = ans = 0;
		for (int i = 1, u, v, l; i < n; i++)
		{
			scanf("%d%d%d", &u, &v, &l);
			addarc(u, v, l);
			addarc(v, u, l);
		}
		dfs(1, 0);
		printf("%d\n", ans);
	}
	return 0;
}