POJ 1741 Tree ——点分治

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define maxn 20005
#define inf 0x3f3f3f3f
using namespace std;

int n,k,h[maxn],to[maxn],ne[maxn],w[maxn],en,cnt;
int a[maxn],b[maxn],ban[maxn],siz[maxn],size,root;
int mx[maxn],now,now_mx,tot;

void init(){en=cnt=0;memset(h,-1,sizeof h);memset(ban,0,sizeof ban);}
void add(int a,int b,int c){to[en]=b;w[en]=c;ne[en]=h[a];h[a]=en++;}
void rd(){for(int i=1;i<n;++i){int a,b,c;scanf("%d%d%d",&a,&b,&c);add(a,b,c);add(b,a,c);}}

void dfs_size(int o,int fa)
{
	siz[o]=1;
	mx[o]=0;
	for (int i=h[o];i>=0;i=ne[i])
	if (to[i]!=fa&&!ban[to[i]]){
		dfs_size(to[i],o);
		siz[o]+=siz[to[i]];
		mx[o]=max(mx[o],siz[to[i]]);
	}
}

void dfs_root(int o,int fa)
{
	int now=max(mx[o],size-siz[o]);
	if (now<now_mx) now_mx=now,root=o;
	for (int i=h[o];i>=0;i=ne[i]) if (to[i]!=fa&&!ban[to[i]]) dfs_root(to[i],o);
}

void dfs_dist(int o,int fa,int dist)
{
	a[++tot]=dist;
	for (int i=h[o];i>=0;i=ne[i])
		if (to[i]!=fa&&!ban[to[i]])
			dfs_dist(to[i],o,dist+w[i]);
}

int cal()
{
	int ret=0,j=tot;
	sort(a+1,a+tot+1);
	for (int i=1;i<=tot;++i)
	{
		while (a[j]+a[i]>k&&j) j--;
		ret+=j;
		if (j>i) ret--;
	}
	return ret/2;
}

void dfs(int o)
{
	dfs_size(o,0);
	size=siz[o];
	now_mx=inf;
	dfs_root(o,0);
	ban[root]=1;
	for (int i=h[root];i>=0;i=ne[i])
	{
		if (!ban[to[i]])
		{
			tot=0;
			dfs_dist(to[i],root,w[i]);
			cnt-=cal();
		}
	}
	tot=0;
	dfs_dist(root,0,0);
	cnt+=cal();
	for (int i=h[root];i>=0;i=ne[i])
		if (!ban[to[i]])
			dfs(to[i]);
}

int main()
{
	while (scanf("%d%d",&n,&k)!=EOF&&n&&k)
	{
		init();rd();dfs(1);
		printf("%d\n",cnt);
	}
}