POJ1741 Tree

http://poj.org/problem?id=1741
题意：

思路：

代码参考这个blog：https://www.cnblogs.com/GXZlegend/p/6641720.html

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<vector>
#include<cstring>
using namespace std;
const int maxn=10000+100;

int n,m,deep[maxn],sz[maxn],f[maxn],vis[maxn],d[maxn],tot,root,sn,ans;
struct Edge{
	int to,dist;
};
vector<Edge> edges;
vector<int> G[maxn];

void AddEdge(int x,int y,int z)
{
	edges.push_back((Edge){y,z});
	G[x].push_back(edges.size()-1);
	edges.push_back((Edge){x,z});
	G[y].push_back(edges.size()-1);
}

void getroot(int u,int fa)
{
	f[u]=0;sz[u]=1;
	for(int i=0;i<G[u].size();i++)
	{
		int v=edges[G[u][i]].to;
		if(v==fa || vis[v])continue;
		getroot(v,u);sz[u]+=sz[v];f[u]=max(f[u],sz[v]);
	}
	f[u]=max(f[u],sn-sz[u]);
	if(f[u]<f[root])root=u;
}

void getdeep(int u,int fa)
{
	d[++tot]=deep[u];
	for(int i=0;i<G[u].size();i++)
	{
		int v=edges[G[u][i]].to,w=edges[G[u][i]].dist;
		if(v==fa || vis[v])continue;
		deep[v]=deep[u]+w;
		getdeep(v,u);
	}
}

int calc(int u)
{
	tot=0;getdeep(u,0);sort(d+1,d+1+tot);
	int l=1,r=tot,sum=0;
	while(l<=tot)
	{
		while(r>l && d[r]+d[l]>m)r--;
		if(l==r)break;
		sum+=r-l;
		l++;
	}
	return sum;
}

void dfs(int u)
{
	deep[u]=0;vis[u]=1;ans+=calc(u);
	for(int i=0;i<G[u].size();i++)
	{
		Edge& e=edges[G[u][i]];
		int v=e.to,w=e.dist;
		if(vis[v])continue;
		deep[v]=w;ans-=calc(v);sn=sz[v];root=0;getroot(v,0);dfs(root);
	}
}

int main()
{
	//freopen("input.in","r",stdin);
	int x,y,z;
	while(cin>>n>>m && n)
	{
		sn=n;
		ans=0;
		memset(vis,0,sizeof(vis));
		for(int i=1;i<=n;i++)G[i].clear();
		edges.clear();
		for(int i=1;i<n;i++)
		{
			scanf("%d%d%d",&x,&y,&z);
			AddEdge(x,y,z);
		}
		root=0;f[0]=(1<<30);
		getroot(1,0);
		dfs(root);
		printf("%d\n",ans);
	}
	return 0;
}