[POJ1741][JZOJ1166] 树中点对距离（点分治模板）

Summary

给出一棵带边权的树，问有多少对点的距离<=Len

Solution

显然应该用点分治。

对于当前递归到的子树的所有路径，我们只考虑经过根的，不经过根的会在以后递归到。

要么是从根出发，要么是根在路径上。

先找出重心。

然后重心每棵子树做，将该子树的所有点到重心的距离存起来，排个序和前面做过的比较计算答案，再加到前面的数组中，所有子树做完以后继续向下分治。

每一层$O(Nlog N)$，有$O(log N)$层，复杂度$O(Nlog^2N)$

Code

#include<cstdio>
#include<cstdlib>
#include<cmath>
#include<cstring>
#include<algorithm>
#include<iostream>
#define fo(i,a,b) for(i=a;i<=b;i++)
#define fod(i,a,b) for(i=a;i>=b;i--)
#define N 10005
using namespace std;
int nt[2*N],dt[2*N],fs[N],lt[N],sz[N],ans,n,g[N],c[N],c1[N],pr[2*N],nm[N],ct,s,l;
bool bz[2*N];
void putin(int p,int x,int y,int z)
{
    dt[p]=y;
    pr[p]=z;
    if (fs[x]==0) fs[x]=p;
    lt[x]=nt[lt[x]]=p;
}
int gets(int k,int f,int s)
{
    int i,sum=1;
    g[++g[0]]=s;
    for(i=fs[k];i>0;i=nt[i])
    {
        int p=dt[i];
        if (p==f||bz[p]) continue;
        sum+=gets(p,k,s+pr[i]);
    }
    return sum;
}
void dg(int k,int f,int num)
{
    int i,v=0;
    sz[k]=1;
    for(i=fs[k];i>0;i=nt[i])
    {
        int p=dt[i];
        if (p==f||bz[p]) continue;
        dg(p,k,num);
        sz[k]+=sz[p];
        v=max(v,sz[p]);
    }
    if(max(num-sz[k],v)<s) ct=k,s=max(num-sz[k],v);
}
void find(int k,int f,int num)
{
    int i,x,y;
    sz[1]=0;
    s=n+1;
    dg(k,0,num); 
    c[0]=0;
    bz[ct]=1;
    nm[0]=0;
    for(i=fs[ct];i>0;i=nt[i])
    {
        int p=dt[i];
        if (p==f||bz[p]) continue;
        g[0]=0;
        nm[++nm[0]]=gets(p,ct,pr[i]);
        sort(g+1,g+g[0]+1);
        x=1;
        y=g[0];
        while (y>0&&(y>0||x<=c[0]))
        {
            if (g[y]<=l) 
            {
                ans++;
                while(x<=c[0]&&g[y]+c[x]<=l) ans+=y,x++;
            }
            y--;
        }
        fo(x,1,c[0]) c1[x]=c[x];
        x=y=1;
        c1[0]=c[0];
        c[0]=0;
        while(x<=c1[0]||(y<=g[0]&&g[y]<=l))
        {
            c[0]++;
            if ((c1[x]>g[y]||c1[0]==0||x>c1[0])&&y<=g[0]&&g[y]<=l) c[c[0]]=g[y++];
            else c[c[0]]=c1[x++]; 
        }
    }
    x=0;
    for(i=fs[ct];i>0;i=nt[i]) if (dt[i]!=f&&!bz[dt[i]]) find(dt[i],ct,nm[++x]);
}
int main()
{
    cin>>n>>l;
    int i,j,k;
    memset(bz,0,sizeof(bz));
    fo(i,1,n-1)
    {
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        putin(2*i-1,x,y,z);
        putin(2*i,y,x,z);
    }
    ans=0;
    find(1,0,n);
    cout<<ans;
}