poj1714(树的分治模板)

解题思路:已经把一些重要的地方注释在代码里了

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int inf=0x3f3f3f3f;
struct node 
{
	int v,w,next;
}s[20003];
int head[10003],cnt,p[10003],f[10003],vis[10003],root,num;//p为以i为根的子树总和,f为以i为根的最大子树 
int r[10003],res,dep[10003],n,k,size;
void sc(int u,int fa)//寻找重心
{
	p[u]=1;f[u]=0;
	for(int i=head[u];i!=-1;i=s[i].next)
	{
		int v=s[i].v;
		if(vis[v]||v==fa) continue;
		sc(v,u);
		p[u]+=p[v];
		f[u]=max(f[u],p[v]);
	}
	f[u]=max(f[u],size-p[u]);
	if(f[u]<f[root]) root=u;
} 
void add_edge(int u,int v,int w)
{
	s[cnt].v=v;
	s[cnt].w=w;
	s[cnt].next=head[u];
	head[u]=cnt++;
	
	s[cnt].v=u;
	s[cnt].w=w;
	s[cnt].next=head[v];
	head[v]=cnt++;	
}
void find_dep(int u,int fa)
{
	r[num++]=dep[u];
	for(int i=head[u];i!=-1;i=s[i].next)
	{
		if(s[i].v==fa||vis[s[i].v]) continue;
		dep[s[i].v]=dep[u]+s[i].w;
		find_dep(s[i].v,u);
	}
}
int find_sum(int u,int d)
{
	int tp=0;
	dep[u]=d;
	num=0;
	find_dep(u,0);
	sort(r,r+num);
	int j=num-1;
	for(int i=0;i<num;i++)
	{
		while(j>i&&r[j]+r[i]>k) j--;
		tp+=max((j-i),0);
	}
	return tp;
}
void find_ans(int u)
{	
	vis[u]=1;
	res+=find_sum(u,0);// 将u为根节点下的所有距离小于k的节点对数都加上 
	for(int i=head[u];i!=-1;i=s[i].next)
	{
		int v=s[i].v;
		if(!vis[v])
		{
			res-=find_sum(v,s[i].w);//减去上面u下面同一子树中的也算的情况 
			size=p[v];
			root=0;
			sc(v,0);
			find_ans(root);
		}
	}
}
int main()
{
	//freopen("t.txt","r",stdin);
	int u,v,w;
	while(scanf("%d%d",&n,&k)!=EOF&&(n+k))
	{
		memset(head,-1,sizeof(head));
		cnt=0;
		for(int i=0;i<n-1;i++)
		{
			scanf("%d%d%d",&u,&v,&w);
			add_edge(u,v,w);
		}
		memset(vis,0,sizeof(vis));
		size=n;
		root=0;
		res=0;
		f[root]=inf;
		sc(1,0);
		find_ans(root);
		printf("%d\n",res);
	}
	return 0;
 } 

猜你喜欢

转载自blog.csdn.net/qq_39861441/article/details/87915937