POJ 1741tree-点分治入门

学习了一下点分治，如果理解有误还请不吝赐教。

为了快速求得树上任意两点之间距离满足某种关系的点对数，我们需要用到这种算法。

点分治是树上的一种分治算法，依靠树和子树之间的关系进行分治从而降低复杂度。

和其他树上的算法有一些区别的是，点分治算法不是先处理局部信息，再将他们汇总得到整个树的信息。点分治处理的问题一般不是树整体的信息，而是树上局部的关系，这就导致我们不能将它看作一个整体，而应该从一开始就处理，在从上往下处理的过程中不断完善信息。在这种思想下我觉得能够更好的理解这个算法。

例如：

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.
Input
The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.
Output
For each test case output the answer on a single line.
Sample Input
5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0
Sample Output
8

题目的大概意思就是说，想要找到树上两点之间距离小于k的点对数（路径是有长度的）。我们如何使用点分治算法来解决这个问题呢？

为了找到合适的分治方法，我们引入一个概念叫做树的重心。树的重心指的是以该节点为根形成的最大子树规模最小的节点。听起来可能有点绕，其实还是挺符合直觉的。指的就是，一个树有多个节点，每个节点都可以作为这个树的根，而一旦根节点确定，就会有多个子树，这些子树中规模最大的一般是根节点下直接相连的某个子树。所谓树的重心，就是这个最大子树规模最小的节点。比如：
在上面这个图中，我们选择i节点作为树的重心，它的子树中规模最大的是3，选择其他的都会比3大。

知道什么是重心以后，我们来看如何求一个树的重心。直接看代码吧

void GetRoot(int x,int father)
{
	int v;
	SonNum[x]=1;//子树节点的总个数
	MaxNum[x]=1;//最大子树的节点个数
	for(int i=Last[x];i;i=Edge[i].last)
	{
		v=Edge[i].to; 
		//Vis[v]是用来标记是否已经处理过该节点了，这个标记会在后面修改，如果已经处理过我们就不要将这个节点再考虑进来了
		if(v==father || Vis[v]) continue;
		GetRoot(v,x);
		SonNum[x]+=SonNum[v];
		if(SonNum[v]>MaxNum[x]) MaxNum[x]=SonNum[x];
	}
	MaxNum[x]=max(MaxNum[x],ss-SonNum[x]);
	if(rootx>MaxNum[x]) root=x,rootx=MaxNum[x];
}

得到树的重心以后我们就要对他进行处理，首先当然是统计子树上任意一点到树根（也就是树的重心）的距离

void GetDis(int x,int father,int dis)
{
	int v;
	//Dis数组记录树上其他点到重心的距离，因为我们不需要知道哪个点，所以直接保存就可以
	Dis[++dlen]=dis;
	for(int i=Last[x];i;i=Edge[i].last)
	{
		v=Edge[i].to; if(v==father|| Vis[v]) continue;
		GetDis(v,x,dis+Edge[i].len);
	}
}

得到距离以后我们就可以根据题目的要求进行计数啦。这里要求的是任意两点的距离小于k，那我们就要先得到两点间的距离。在以重心为根的树上，在两个不同子树上的点的距离就是他们距离重心的距离和。可是如果在同一个子树上的话就不是这样了。但是我们不好确定两个节点是否在同一个子树上，所以先囫囵吞枣将同一个子树上的都计算上，然后再访问子树将他们减去。

int Count(int x,int dis)
{
	for(int i=0;i<=dlen;++i) Dis[i]=0;
	dlen=0;
	GetDis(x,0,dis);
	sort(Dis+1,Dis+1+dlen);
	int l=1,r=dlen,ret=0;
	//如果l到r的距离小于k，则l到l和r之间的任意一点的距离都小于k，所以直接加上r-l
	while(l<=r)
	{
		if(Dis[l]+Dis[r]<=kk) ret+=r-l,l++;
		else r--;
	}
	return ret;
}

但是显然这样是多算的，我们要减去在同一个子树上的满足条件的节点，他们会在更后面再次加上。

void Solve(int x)
{
	int v;
	ans+=Count(x,0);
	Vis[x]=true;
	for(int i=Last[x];i;i=Edge[i].last)
	{
		v=Edge[i].to; if(Vis[v]) continue;
		//在这里减去同一个子树的上错误加上的点
		ans-=Count(v,Edge[i].len);
		ss=SonNum[v]; rootx=INT_MAX; root=0;
		//处理子树，再加上正确的点
		GetRoot(v,x);
		Solve(root);
	}
}

可能稍微有些难以理解的是为什么这样就可以减去刚开始错误地加上的同一个子树上的点。这里稍微解释一下：

还是以这个图为例，假如我们一开始处理的是树的重心i节点，那么我们正确计算的就是分布在四个子树上之间的距离，错误计算的就是同一个子树之间的距离，例如：D-A-i-A-E和E-A-i-A等，为了处理这个问题，我们后面又访问了一下子节点，注意上面的Solve函数中的ans-=Count(v,Edge[i].len);，为什么这样写就可以减去子节点的影响呢？需要注意我们已经将重心的Vis[x]的值已经修改，因此子节点无法访问除了当前子树外的其他子树，而且有一个初值Edge[i].len，这样处理以后他们的Dis数组的值和之前从重心访问是相同的。也就是说，之前会计入答案的，这里也会再次记入答案，且只计入了同一个子树中的。减去他们以后就是正确的数目。

然后我们再访问子树。将子树看作一个单独的树，再次同样的处理。

AC代码：

#include<iostream>
#include<cstring>
#include<cstdio>
#include<climits>
#include<algorithm>
#include<ctime>
#include<cstdlib>
#include<queue>
#include<set>
#include<map>
#include<cmath>

using namespace std;

const int MAXN=1e5+5;
struct edge
{
	int to,len,last;
}Edge[MAXN<<1]; int Last[MAXN],tot;
int n,kk,SonNum[MAXN],MaxNum[MAXN],Vis[MAXN],Dis[MAXN];
int ans,root,rootx,dlen,ss;

int getint()
{
	int x=0,sign=1; char c=getchar();
	while(c<'0' || c>'9')
	{
		if(c=='-') sign=-1; c=getchar();
	}
	while(c>='0' && c<='9')
	{
		x=x*10+c-'0'; c=getchar();
	}
	return x*sign;
}

void Init()
{
	for(int i=0;i<=n;++i) Last[i]=0; tot=0; 
	ans=0; for(int i=0;i<=n;++i) Vis[i]=false;
}

void AddEdge(int u,int v,int w)
{
	Edge[++tot].to=v; Edge[tot].len=w; 
	Edge[tot].last=Last[u]; Last[u]=tot;
}

void Read()
{
	int u,v,w;
	for(int i=1;i<n;i++)
	{
		u=getint(); v=getint(); w=getint();
		AddEdge(u,v,w); AddEdge(v,u,w);
	}
}

void GetRoot(int x,int father)
{
	int v;
	SonNum[x]=1; MaxNum[x]=1;
	for(int i=Last[x];i;i=Edge[i].last)
	{
		v=Edge[i].to; if(v==father || Vis[v]) continue;
		GetRoot(v,x);
		SonNum[x]+=SonNum[v];
		if(SonNum[v]>MaxNum[x]) MaxNum[x]=SonNum[x];
	}
	MaxNum[x]=max(MaxNum[x],ss-SonNum[x]);
	if(rootx>MaxNum[x]) root=x,rootx=MaxNum[x];
}

void GetDis(int x,int father,int dis)
{
	int v;
	Dis[++dlen]=dis;
	for(int i=Last[x];i;i=Edge[i].last)
	{
		v=Edge[i].to; if(v==father|| Vis[v]) continue;
		GetDis(v,x,dis+Edge[i].len);
	}
}

int Count(int x,int dis)
{
	for(int i=0;i<=dlen;++i) Dis[i]=0;
	dlen=0;
	GetDis(x,0,dis);
	sort(Dis+1,Dis+1+dlen);
	int l=1,r=dlen,ret=0;
	while(l<=r)
	{
		if(Dis[l]+Dis[r]<=kk) ret+=r-l,l++;
		else r--;
	}
	return ret;
}

void Solve(int x)
{
	int v;
	ans+=Count(x,0);
	Vis[x]=true;
	for(int i=Last[x];i;i=Edge[i].last)
	{
		v=Edge[i].to; if(Vis[v]) continue;
		ans-=Count(v,Edge[i].len);
		ss=SonNum[v]; rootx=INT_MAX; root=0;
		GetRoot(v,x);
		Solve(root);
	}
}

void Work()
{
	rootx=INT_MAX; ss=n; root=0;
	GetRoot(1,0); 
	Solve(root);
}

void Write()
{
	printf("%d\n",ans); 
}

int main()
{
	while(1)
	{
		Init();
		n=getint(); kk=getint();
		if(n==0 && kk==0) break;
		Read();
		Work();
		Write();
	}
	return 0;
}

参考博客：传送门