POJ 1741 Tree 树的点分治

点击打开链接

题意:给出n个结点的树,边的权值,k<=1e9,n<=1e4,求dist(u,v)<=k的个数?
dist(u,v)要么经过根结点,要么就在某个子树中,则分治来做

设d[i]为i到当前根的距离,通过rt的合法对数为d[u]+d[v]<=k && u,v不在同一个子树中
不通过rt的路径为子问题,递归求解即可

可以用反证法 证明若选取树的重心为根,则最大子树的结点数<=n/2,所以最多log(n)层
每层都做了排序,复杂度为O(nlogn),总的时间复杂度为O(nlog^2n) 

 

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<int,int> ii;
const int N=2e5+20;
const ll mod=1e9+7; 
int n,k,ans,d[N];//d[i] i->rt
vector<ii> e[N]; 
bool done[N];//
int root,f[N],s[N],size;
vector<int> dep;//
void getroot(int u,int fa)//选重心,做根 
{
	s[u]=1,f[u]=0;//结点数,和最大子树的结点数
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i].first;
		if(v==fa||done[v]) continue;
		getroot(v,u);
		s[u]+=s[v];	
		f[u]=max(f[u],s[v]);
	}
	f[u]=max(f[u],size-s[u]);//以u为根的最大子树的结点数
	if(f[u]<f[root])
		root=u; 
}
void getdep(int u,int fa)
{
	dep.push_back(d[u]);
	s[u]=1;
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i].first,w=e[u][i].second;
		if(v==fa||done[v]) continue;
		
		d[v]=d[u]+w;
		getdep(v,u);
		s[u]+=s[v];
	}
}
int calc(int u,int init)
{
	dep.clear(),d[u]=init;
	getdep(u,0); 
	sort(dep.begin(),dep.end());//排序后求 di+dj<=k的对数 
	int res=0;
	for(int l=0,r=dep.size()-1;l<r;)//枚举l 找到最大r
	{
		if(dep[l]+dep[r]<=k)
		{
			res+=r-l;
			l++;
		}
		else
		r--;
	}
	return res;
}
void work(int u)
{
	ans+=calc(u,0);//经过根u的合法路径 
	done[u]=true;
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i].first,w=e[u][i].second;
		if(done[v]) continue;
		ans-=calc(v,w);//计算经过根u时,扣除掉来自同一个子树 
		f[0]=size=s[u];// 
		getroot(v,root=0);
		work(root);
	}
}
int main()
{
	while(cin>>n>>k&&(n+k))
	{
		int u,v,w;	
		for(int i=1;i<=n;i++)
		{
			done[i]=false;
			e[i].clear();
		}
		for(int i=1;i<=n-1;i++)
		{
			scanf("%d%d%d",&u,&v,&w);
			e[u].push_back(ii(v,w));
			e[v].push_back(ii(u,w));
		}
		f[0]=size=n;//
		getroot(1,root=0);
		ans=0;
		work(root);
		printf("%d\n",ans);
	} 
	return 0;
}