【树上点分治1】Tree POJ 1741(容斥原理)

【树上点分治1】Tree POJ 1741(容斥原理)

学习博客

题意：
给定一棵树，要求找出树上两点之间距离不超过k的点对有几个。

思路：

以A为根节点，把树上路径分成两部分

1.经过A节点

用tmp[]存储从A节点到各节点的距离，那么ans=tmp[i]+tm[j]<=k，但是若i，j在同一颗子树，比如A->B->D和A->B->E，则不能进行合并，所有需要减去在同一颗子树的符合情况的点对，这刚好可以对B子树进行相同的处理，以B为根节点算一次ans’，再用A节点的ans减去即可

2.不经过A节点

递归子树处理

#include<iostream>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
const int inf=2e7+5;
const int maxn=2e5+5;
int n,m,k1,k,rt,cnt,head[maxn],to[maxn],nex[maxn],w[maxn];
int maxp[maxn],tmp[maxn],judge[maxn],siz[maxn],vis[maxn];
int ans,dis[maxn],q[maxn],que[maxn];
void add(int x,int y,int z){
	nex[++k1]=head[x];
	to[k1]=y;
	w[k1]=z;
	head[x]=k1;
}
void getzx(int x,int f){//dp求重心 
	siz[x]=1,maxp[x]=0;
	for(int i=head[x];i;i=nex[i]){
		int y=to[i];
		if(y==f||vis[y])continue;
		getzx(y,x);
		siz[x]+=siz[y];
		maxp[x]=max(maxp[x],siz[y]);
	}
	maxp[x]=max(maxp[x],cnt-siz[x]);
	if(maxp[x]<maxp[rt])rt=x;
}
void getdis(int x,int f){
	tmp[++tmp[0]]=dis[x];
	for(int i=head[x];i;i=nex[i]){
		int y=to[i];
		if(y==f||vis[y])continue;
		dis[y]=dis[x]+w[i];
		getdis(y,x);
	}
}

int calc(int x,int val){
	int sum=0;
	dis[x]=val,tmp[0]=0;//初始化
	getdis(x,0);
	sort(tmp+1,tmp+1+tmp[0]);
	int l=1,r=tmp[0];
	while(l<r){
		if(tmp[l]+tmp[r]>k)r--;
		else sum+=r-l,l++;
	}
	return sum;
}
void solve(int x){
	vis[x]=1;
	ans+=calc(x,0);//方便计算
	for(int i=head[x];i;i=nex[i]){
		int y=to[i];
		if(vis[y])continue;
		ans-=calc(y,w[i]);//容斥原理 ：减去重复计数 
		cnt=siz[y],maxp[rt=0]=inf;//初始化 
		getzx(y,0);//找子树重心 
		solve(rt);//对重心 
	}
}
int main(){
	while(scanf("%d%d",&n,&k)){
		if(n==0&&k==0)break;
	    ans=k1=0;
	    memset(head,0,sizeof(head));
	    memset(vis,0,sizeof(vis));
	    for(int i=1,x,y,z;i<n;i++){
		    scanf("%d%d%d",&x,&y,&z);
		    add(x,y,z),add(y,x,z);
	    }
	    cnt=n;
	    maxp[rt=0]=inf;
	    getzx(1,0);
	    solve(rt);	
	    printf("%d\n",ans);
	}
}