POJ 1987 Distance Statistics

http://poj.org/problem?id=1987

题意:给一棵树，求树上有多少对节点满足距离<=K

思路:点分治，我们考虑把每个距离都存起来，然后排序，一遍扫描计算一下，注意还要减掉自己加自己的方案。而且，我们还要去掉走到同一个子树的方案。复杂度:O(nlog^2n)

#include<cstdio>
#include<cmath>
#include<cstring>
#include<iostream>
#include<algorithm>
#define ll long long
int tot,go[1000005],first[1000005],next[1000005];
ll st[1000005],val[1000005];
int sum,son[1000005],root,n,F[1000005],c[1000005];
int pd[1000005],sz,vis[1000005];
ll dis[1000005];
int cnt,K,ans;
int read(){
int t=0,f=1;char ch=getchar();
while (ch<'0'||ch>'9'){if (ch=='-') f=-1;ch=getchar();}
while ('0'<=ch&&ch<='9'){t=t*10+ch-'0';ch=getchar();}
return t*f;
}
void insert(int x,int y,int z){
tot++;
go[tot]=y;
next[tot]=first[x];
first[x]=tot;
val[tot]=z;
}
void add(int x,int y,int z){
insert(x,y,z);insert(y,x,z);
}
void findroot(int x,int fa){
son[x]=1;F[x]=0;
for (int i=first[x];i;i=next[i]){
int pur=go[i];
if (pur==fa||vis[pur]) continue;
findroot(pur,x);
son[x]+=son[pur];
F[x]=std::max(F[x],son[pur]);
}
F[x]=std::max(F[x],sum-son[x]);
if (F[x]<F[root]) root=x;
}
void bfs(int x){
int h=1,t=1;c[1]=x;pd[x]=sz;dis[x]=0;
while (h<=t){
int now=c[h++];
for (int i=first[now];i;i=next[i]){
int pur=go[i];
if (vis[pur]||pd[pur]==sz) continue;
pd[pur]=sz;
dis[pur]=dis[now]+val[i];
c[++t]=pur;
st[++cnt]=dis[pur];
}
}
std::sort(st+1,st+1+cnt);
int j=cnt,res=0,Cnt=0;
for (int i=1;i<=t;i++){
while (j>1&&st[i]+st[j]>K) j--;
if (st[i]+st[j]<=K) res+=j;
if (st[i]+st[i]<=K) Cnt++;
}
res-=Cnt;
ans+=res/2;
}
int del(int x,int Dis){
dis[x]=Dis;sz++;
int h=1,t=1;cnt=1;st[cnt]=Dis;
pd[x]=sz;c[1]=x;
while (h<=t){
int now=c[h++];
for (int i=first[now];i;i=next[i]){
int pur=go[i];
if (pd[pur]==sz||vis[pur]) continue;
dis[pur]=dis[now]+val[i];
st[++cnt]=dis[pur];
pd[pur]=sz;
c[++t]=pur;
}
}
int j=cnt,res=0,Cnt=0;
std::sort(st+1,st+1+cnt);
for (int i=1;i<=t;i++){
while (j>1&&st[i]+st[j]>K) j--;
if (st[i]+st[j]<=K) res+=j;
if (st[i]+st[i]<=K) Cnt++;
}
res-=Cnt;
return res/2;
}
void solve(int x){
vis[x]=1;++sz;
cnt=1;st[cnt]=0;
bfs(x);
for (int i=first[x];i;i=next[i]){
int pur=go[i];
if (vis[pur]) continue;
ans-=del(pur,val[i]);
}
int Cnt=sum;
for (int i=first[x];i;i=next[i]){
int pur=go[i];
if (vis[pur]) continue;
if (son[pur]>son[x]) sum=Cnt-son[x];
else sum=son[pur];
root=0;
findroot(pur,x);
solve(root);
}
}
int main(){
int m;
char s[20];
scanf("%d%d\n",&n,&m);
for (int i=1;i<n;i++){
int x,y,z;
scanf("%d%d%d",&x,&y,&z);
add(x,y,z);
scanf("%s",s+1);
}
scanf("%d\n",&K);
F[0]=0x7fffffff;
root=0;sum=n;
findroot(1,0);
solve(root);
printf("%d\n",ans);
}