poj1741_Tree(树的点分治入门题)

题目链接：poj1741_Tree

题意：

给你一颗n个节点的树，每条边有一个值，问有多少点对(u,v)，满足u->v的最短路径小于k。

题解：

典型的树的分治，板子题。

#include<cstdio>
#include<algorithm>
#include<cstring>
#define F(i,a,b) for(int i=a;i<=b;++i)
using namespace std;

const int N=1e4+7;

int n,k,g[N],v[N*2],nxt[N*2],w[N*2],ed,ans;
int vis[N],size[N],mx[N],mi,dis[N],tot,root;

inline void adg(int x,int y,int z){v[++ed]=y,w[ed]=z,nxt[ed]=g[x],g[x]=ed;}
void init(){F(i,1,n)g[i]=0,vis[i]=0;ed=ans=0;}

void dfs_size(int u,int fa)
{
    size[u]=1,mx[u]=0;
    for(int i=g[u];i;i=nxt[i])
        if(v[i]!=fa&&!vis[v[i]])
        {
            dfs_size(v[i],u),size[u]+=size[v[i]];
            if(size[v[i]]>mx[u])mx[u]=size[v[i]];
        }
}

void dfs_root(int r,int u,int fa)
{
    if(size[r]-size[u]>mx[u])mx[u]=size[r]-size[u];
    if(mx[u]<mi)mi=mx[u],root=u;
    for(int i=g[u];i;i=nxt[i])
        if(v[i]!=fa&&!vis[v[i]])
            dfs_root(r,v[i],u);
}

void dfs_dis(int u,int d,int fa)
{
    dis[++tot]=d;
    for(int i=g[u];i;i=nxt[i])
        if(v[i]!=fa&&!vis[v[i]])
            dfs_dis(v[i],d+w[i],u);
}

int calc(int u,int d)
{
    int ans=0;
    tot=0,dfs_dis(u,d,0);
    sort(dis+1,dis+1+tot);
    int i=1,j=tot;
    while(i<j)
    {
        while(dis[i]+dis[j]>k&&i<j)j--;
        ans+=j-i,i++;
    }
    return ans;
}

void dfs(int u=1)
{
    mi=n,dfs_size(u,0);
    dfs_root(u,u,0);
    ans+=calc(root,0),vis[root]=1;
    for(int i=g[root];i;i=nxt[i])
        if(!vis[v[i]])ans-=calc(v[i],w[i]),dfs(v[i]);
}

int main()
{
    while(~scanf("%d%d",&n,&k)&&n+k!=0)
    {
        init();
        F(i,1,n-1)
        {
            int x,y,z;
            scanf("%d%d%d",&x,&y,&z);
            adg(x,y,z),adg(y,x,z);
        }
        dfs();
        printf("%d
",ans);
    }
    return 0;
}