【poj1987】 Distance Statistics

http://poj.org/problem?id=1987 (题目链接)

题意

　　给出一棵树，求树上距离不超过K的点对个数。

Solution

　　点分治，同poj1741。

代码

// poj1987
#include<algorithm>
#include<iostream>
#include<cstdlib>
#include<cstring>
#include<cstdio>
#include<cmath>
#define LL long long
#define MOD 100000000
#define inf 2147483640
#define Pi acos(-1.0)
#define free(a) freopen(a".in","r",stdin),freopen(a".out","w",stdout);
using namespace std;

const int maxn=40010;
struct edge {int to,next,w;}e[maxn<<1];
int head[maxn],deep[maxn],d[maxn],size[maxn],f[maxn],vis[maxn];
int n,m,K,cnt,rt,sum,ans;

void link(int u,int v,int w) {
	e[++cnt].to=v;e[cnt].next=head[u];head[u]=cnt;e[cnt].w=w;
	e[++cnt].to=u;e[cnt].next=head[v];head[v]=cnt;e[cnt].w=w;
}
void calroot(int x,int fa) {
	size[x]=1;f[x]=0;
	for (int i=head[x];i;i=e[i].next)
		if (!vis[e[i].to] && e[i].to!=fa) {
			calroot(e[i].to,x);
			size[x]+=size[e[i].to];
			f[x]=max(f[x],size[e[i].to]);
		}
	f[x]=max(f[x],sum-size[x]);
	if (f[x]<f[rt]) rt=x;
}
void caldeep(int x,int fa) {
	deep[++deep[0]]=d[x];
	for (int i=head[x];i;i=e[i].next)
		if (!vis[e[i].to] && e[i].to!=fa) {
			d[e[i].to]=d[x]+e[i].w;
			caldeep(e[i].to,x);
		}
}
int cal(int x,int now) {
	d[x]=now;deep[0]=0;
	caldeep(x,0);
	sort(deep+1,deep+deep[0]+1);
	int t=0;
	for (int l=1,r=deep[0];l<r;) {
		if (deep[l]+deep[r]<=K) t+=r-l,l++;
		else r--;
	}
	return t;
}
void solve(int x) {
	vis[x]=1;
	ans+=cal(x,0);
	for (int i=head[x];i;i=e[i].next) if (!vis[e[i].to]) ans-=cal(e[i].to,e[i].w);
	for (int i=head[x];i;i=e[i].next) if (!vis[e[i].to]) {
			rt=0;sum=size[e[i].to];
			calroot(e[i].to,0);
			solve(rt);
		}
}
int main() {
	scanf("%d%d",&n,&m);
	char ch;
	for (int u,v,w,i=1;i<=m;i++) {
		scanf("%d%d%d %c",&u,&v,&w,&ch);
		if (u==v) continue;
		link(u,v,w);
	}
	scanf("%d",&K);
	rt=0;sum=n;f[0]=inf;
	calroot(1,0);
	solve(rt);
	printf("%d",ans);
	return 0;
}