POJ-1741Tree（树分治模板）

题意：求树上距离小于等于m的点对有多少个。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
const int maxn=1e4+100;
struct node{
    int x,w;
    node(int _x,int _w)
    {
        x=_x;
        w=_w;
    }
};
vector<node>a[maxn];
int hvy,n,m,vis[maxn],siz[maxn],ma[maxn];
int S;  //树的大小
void dfs_hvy(int u,int fa)//找重心
{
    siz[u]=1;
    ma[u]=0;
    for(int i=0;i<a[u].size();i++)
    {
        int v=a[u][i].x;
        if(v==fa||vis[v]) continue;
        dfs_hvy(v,u);
        siz[u]+=siz[v];
        ma[u]=max(ma[u],siz[v]);
    }
    ma[u]=max(ma[u],S-siz[u]);
    if(!hvy||ma[u]<ma[hvy])// 找重心，以重心为根的树的最大子树最小
        hvy=u;
}

int dep[maxn],tmp[maxn];
void dfs_dep(int u,int fa)
{
    tmp[++tmp[0]]=dep[u];
    for(int i=0;i<a[u].size();i++)
    {
        int v=a[u][i].x,w=a[u][i].w;
        if(v==fa||vis[v]) continue;///注意找的是未访问的
        dep[v]=dep[u]+w;
        dfs_dep(v,u);
    }
}
int dfs_sum(int u,int d)
{
    tmp[0]=0;
    int res=0;
    dep[u]=d;
    dfs_dep(u,0);//找出孩子节点到u的距离
    sort(tmp+1,tmp+tmp[0]+1);///
    int j=1;
    for(int i=1,j=tmp[0];i<=tmp[0];i++)
    {
        while(j>i&&tmp[i]+tmp[j]>m)j--;///
        res+=max(j-i,0);///
    }
    return res;
}
int ans;
void dfs_ans(int u)
{
    vis[u]=1;
    ans+=dfs_sum(u,0);//加上经过重心的点对数（加的里面包括了不经过的，需要在下面再减去）
    for(int i=0;i<a[u].size();i++)
    {
        int v=a[u][i].x,w=a[u][i].w;
        if(vis[v]) continue;
        ans-=dfs_sum(v,w);/////不经过的在直接孩子处减去
        S=siz[v],hvy=0;
        dfs_hvy(v,u);
        dfs_ans(hvy);
    }
}
int main()
{
    while(scanf("%d%d",&n,&m),n+m)
    {
        for(int i=1;i<=n;i++)
            a[i].clear();
        memset(vis,0,sizeof vis);
        for(int i=1;i<n;i++)
        {
            int x,y,w;
            scanf("%d%d%d",&x,&y,&w);
            a[x].push_back(node(y,w));
            a[y].push_back(node(x,w));
        }
        S=n;hvy=0;
        dfs_hvy(1,0);//找重心
        ans=0;
        dfs_ans(hvy);
        printf("%d\n",ans);
    }
    return 0;;;
}

从这篇博客学到的