2017.7.18 树上距离=k 思考记录

利用树分治的原理，再加上扫描线的o（nm）更新，可以求出树上任意两点距离==k的点对个数

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define N 20005
#define inf 1e9+7
int hou[N],xia[N],zhong[N],v[N],tot,max1,sz[N],i,rt,dis[N],q[N],n,m,w[N],x,y,z,ans[N];
bool vis[N];
void jia(int x,int y,int z)
{
	++tot,hou[tot]=xia[x],zhong[tot]=y,xia[x]=tot,v[tot]=z;
	//cout<<tot<<" "<<x<<" "<<y<<endl;
}
void jian(int x,int y,int z)
{
	jia(x,y,z);
	jia(y,x,z);
}
int zzx(int o,int fu)
{
sz[o]=1;
int maxx=0;	
for(int i=xia[o];i!=-1;i=hou[i])	
{
	int nd=zhong[i];
	if(nd==fu||vis[nd])continue;
	 zzx(nd,o);
	 sz[o]+=sz[nd];
     maxx=max(maxx,sz[nd]);
}
		 maxx=max(maxx,tot-maxx);
if(maxx<max1){max1=maxx,rt=o;}
}
void dfs(int o,int fu,int d)
{
	dis[o]=d;
	q[++q[0]]=d;
	for(int i=xia[o];i!=-1;i=hou[i])
	{
	int nd=zhong[i];
	//cout<<o<<" "<<nd<<endl;
	if(nd==fu||vis[nd])continue;
	dfs(nd,o,d+v[i]);		
 	}		
 	
}
void tj1(int o,int fu)
{
	q[0]=0;
	dfs(o,fu,0);
    sort(q+1,q+1+q[0]);
    for(int i=1;i<=m;i++)
    {
    int l=1;
    int r=q[0];
    while(l<r)
    {
    	if(q[l]+q[r]==w[i]){ans[i]++;if(q[l]+q[r-1]==w[i])--r;else ++l;}
    	else if(q[l]+q[r]>w[i])--r;
    	else ++l;
	}
}
}
void tj2(int o,int fu)
{
	q[0]=0;
	dfs(o,fu,dis[o]-dis[fu]);
	  sort(q+1,q+1+q[0]);
    for(int i=1;i<=m;i++)
    {
    int l=1;
    int r=q[0];
    while(l<r)
    {
    	if(q[l]+q[r]==w[i]){ans[i]--;if(q[l]+q[r-1]==w[i])--r;else ++l;}
    	else if(q[l]+q[r]>w[i])++l;
    	else --r;
	}	
}
}
void work(int o,int fu)
{
vis[o]=1;
   tj1(o,fu);
	for(int i=xia[o];i!=-1;i=hou[i])
	{
		int nd=zhong[i];
		if(nd==fu)continue;
		tj2(nd,o);
       	tot=sz[nd];
       	max1=inf;
		zzx(nd,o);	 
		work(rt,o);	
	}
}
int main()
{
	memset(xia,-1,sizeof(xia));
	scanf("%d%d",&n,&m);
	for(i=1;i<n;i++)
	{
		scanf("%d%d%d",&x,&y,&z);
		jian(x,y,z);
	}
	for(i=1;i<=m;i++)
	{
		scanf("%d",&w[i]);
	}
	tot=n;
	max1=inf;
	zzx(1,0);
	work(rt,0);
	for(i=1;i<=m;i++)
	{
		printf("%d ",ans[i]);	
	}
}