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[],first[],next[];
ll st[],val[];
int sum,son[],root,n,F[],c[];
int pd[],sz,vis[];
ll dis[];
int cnt,K,ans;
int read(){
    int t=,f=;char ch=getchar();
    while (ch<''||ch>''){if (ch=='-') f=-;ch=getchar();}
    while (''<=ch&&ch<=''){t=t*+ch-'';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]=;F[x]=;
    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=,t=;c[]=x;pd[x]=sz;dis[x]=;
    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+,st++cnt);
    int j=cnt,res=,Cnt=;
    for (int i=;i<=t;i++){
        while (j>&&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/;
}
int del(int x,int Dis){
    dis[x]=Dis;sz++;
    int h=,t=;cnt=;st[cnt]=Dis;
    pd[x]=sz;c[]=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=,Cnt=;
    std::sort(st+,st++cnt);
    for (int i=;i<=t;i++){
       while (j>&&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/;
}
void solve(int x){
    vis[x]=;++sz;
    cnt=;st[cnt]=;
    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=;
         findroot(pur,x);
         solve(root);
    }
}
int main(){
    int m;
    char s[];
    scanf("%d%d\n",&n,&m);
    for (int i=;i<n;i++){
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);
        scanf("%s",s+);
    }
    scanf("%d\n",&K);
    F[]=0x7fffffff;
    root=;sum=n;
    findroot(,);
    solve(root);
    printf("%d\n",ans);
}