POJ1741 Tree 点分治

问题描述

POJ1741

---

题解

题意是求树上满足两点之间距离 \(dis \le k\) 的点对 \((x,y)\) 的数目。

点分治。

---

\(\mathrm{Code}\)

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>

using namespace std;

template <typename Tp>
void read(Tp &x){
	x=0;char ch=1;int fh;
	while(ch!='-'&&(ch>'9'||ch<'0')) ch=getchar();
	if(ch=='-') ch=getchar(),fh=-1;
	else fh=1;
	while(ch>='0'&&ch<='9') x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
	x*=fh;
}

const int maxn=10007;
const int maxm=20007;
const int INF=0x3f3f3f3f;

int n,k;
int Head[maxn],to[maxm],Next[maxm],w[maxm],tot;
int size[maxn],mx,sz,root;

int ans,dis[maxn];
int q[maxn],l,r;

bool del[maxn];

void reset(){
	memset(Head,0,sizeof(Head));
	memset(Next,0,sizeof(Next));
	memset(del,0,sizeof(del));
	tot=ans=0;
}

void add(int x,int y,int z){
	to[++tot]=y,Next[tot]=Head[x],Head[x]=tot,w[tot]=z;
}

void Init(){
	read(k);
	for(int i=1,x,y,z;i<n;i++){
		read(x);read(y);read(z);
		add(x,y,z);add(y,x,z);
	}
}

void getdis(int x,int fa){
	q[++r]=dis[x];
	for(int i=Head[x];i;i=Next[i]){
		int y=to[i];
		if(y==fa||del[y]) continue;
		dis[y]=dis[x]+w[i];
		getdis(y,x);
	}
}

int calc(int x,int val){
	r=0;dis[x]=val;
	getdis(x,0);
	int res=0;l=1;
	sort(q+1,q+r+1);
	while(l<r){
		if(q[l]+q[r]<=k) res+=r-l,++l;
		else --r;
	}
	return res;
}

void getroot(int x,int fa){
	size[x]=1;int num=0;
	for(int i=Head[x];i;i=Next[i]){
		int y=to[i];
		if(y==fa||del[y]) continue;
		getroot(y,x);
		size[x]+=size[y];
		num=max(num,size[y]);
	}
	num=max(num,sz-size[x]);
	if(num<mx) mx=num,root=x;
}

void dfs(int x){
	ans+=calc(x,0);del[x]=1;
	for(int i=Head[x];i;i=Next[i]){
		int y=to[i];
		if(del[y]) continue;
		ans-=calc(y,w[i]);
		sz=size[y],mx=INF;
		getroot(y,0);dfs(root);
	}
}

void solve(){
	sz=n,mx=INF;
	getroot(1,0);
	dfs(root);
	printf("%d\n",ans);
}

int main(){
	while(1){
		read(n);if(!n) return 0;
		reset();Init();solve();
	}
	return 0;
}