【点分治】poj1741 Tree / poj2114 Boatherds / poj1987 Distance Statistics

三道题都很类似。给出1741的代码

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
#define MAXN 10001
typedef pair<int,int> Point;
int n,K,ans;
int v[MAXN<<1],w[MAXN<<1],first[MAXN],next[MAXN<<1],en;
void AddEdge(const int &U,const int &V,const int &W)
{
	v[++en]=V;
	w[en]=W;
	next[en]=first[U];
	first[U]=en;
}
bool centroid[MAXN];//顶点是否已经作为重心删除的标记
int size[MAXN];//以该顶点为根的子树的大小
//计算子树的大小
int calc_sizes(int U,int Fa)
{
	int res=1;
	for(int i=first[U];i;i=next[i])
	  if(v[i]!=Fa&&(!centroid[v[i]]))
		res+=calc_sizes(v[i],U);
	return size[U]=res;
}
//查找重心的递归函数，nn是整个子树的大小
//在以U为根的子树中寻找一个顶点，使得删除该顶点后得到的最大子树的顶点数最少
//返回值为(最大子树的顶点数,顶点编号)
Point calc_centroid(int U,int Fa,int nn)
{
	Point res=make_pair(2147483647,-1);
	int sum=1,maxv=0;
	for(int i=first[U];i;i=next[i])
	  if(v[i]!=Fa&&(!centroid[v[i]]))
	    {
	      res=min(res,calc_centroid(v[i],U,nn));
	      maxv=max(maxv,size[v[i]]);
	      sum+=size[v[i]];
	    }
	maxv=max(maxv,nn-sum);
	res=min(res,make_pair(maxv,U));
	return res;
}
int td[MAXN],en2,ds[MAXN],en3;
//计算子树中所有顶点到重心的距离的递归函数
void calc_dis(int U,int Fa,int d)
{
	td[en2++]=d;
	for(int i=first[U];i;i=next[i])
	  if(v[i]!=Fa&&(!centroid[v[i]]))
	    calc_dis(v[i],U,d+w[i]);
}
int calc_pairs(int dis[],int En)
{
	int res=0;
	sort(dis,dis+En);
	for(int i=0;i<En;++i)
	  res+=upper_bound(dis+i+1,dis+En,K-dis[i])-(dis+i+1);
	return res;
}
void solve(int U)
{
	calc_sizes(U,-1);
	int s=calc_centroid(U,-1,size[U]).second;
	centroid[s]=1;
	//情况1:递归统计按重心s分割后的子树中的对数
	for(int i=first[s];i;i=next[i])
	  if(!centroid[v[i]])
	    solve(v[i]);
	//情况2:统计经过重心s的对数
	en3=0; ds[en3++]=0;
	for(int i=first[s];i;i=next[i])
	  if(!centroid[v[i]])
	    {
	      en2=0; calc_dis(v[i],s,w[i]);
	      ans-=calc_pairs(td,en2);//先把重复统计的部分(即情况1)减掉
	      memcpy(ds+en3,td,en2*sizeof(int)); en3+=en2;
	    }
	ans+=calc_pairs(ds,en3);
	centroid[s]=0;
}
void init()
{
	memset(first,0,sizeof(first));
	en=ans=0;
}
int main()
{
	while(1)
	  {
	  	scanf("%d%d",&n,&K);
		if(!n&&!K) break;
		init();
		int a,b,c;
		for(int i=1;i<n;++i)
	  	  {
	  		scanf("%d%d%d",&a,&b,&c);
	  		AddEdge(a,b,c); AddEdge(b,a,c);
	 	   }
	 	solve(1);
	 	printf("%d\n",ans);
	  }
	return 0;
}