Tree

Tree

#include<cstdio>
#include<iostream>
#include<cmath>
#include<algorithm>
using namespace std;
const int maxn=100400;
const int inf=0x3f3f3f3f;
struct node
{
    int v,w,next;
} e[maxn];

int num,head[maxn],siz[maxn],f[maxn],dis[maxn],cnt,mins,root,k,ans,n;


/*
    程序中部分变量说明
    dis[i]  所有路径到 重心 的长度
    siz[i]  以 i 为根节点的子树的大小(当前重心下)
    f[i]    i 是否还存在(分治完一次后就把重心删掉了)
    cnt     记录 dis 的个数(即路径个数)
    root    当前子树的重心
    maxs    当前讨论的点所有子树大小中最大值(并不是全局变量，是尝试每个重心时重新开的一个变量)
    mins    讨论过的点的子树大小中最大的最小值(是全局变量，用来确定哪个才是重心)
*/


void add(int u,int v,int w)
{
    num++;
    e[num].v=v;
    e[num].w=w;
    e[num].next=head[u];
    head[u]=num;
}

int get_size(int x,int fa)  //返回以 x 为根的子树大小，其中 x 父节点为 fa
{
    siz[x]=1;
    for (int i=head[x]; i; i=e[i].next)
    {
        int v=e[i].v;
        if (v!=fa&&!f[v])
        {
            siz[x]+=get_size(v,x);
        }
    }
    return siz[x];
}

void get_dis(int x,int d,int fa) //返回以 x 为根的子树大小，其中 x 父节点为 fa
{
    dis[++cnt]=d;
    for (int i=head[x]; i; i=e[i].next)
    {
        int v=e[i].v;
        if (v!=fa&&!f[v])
        {
            get_dis(v,d+e[i].w,x);
        }
    }
}

void dfs_root(int x,int tot,int fa)
{
    //求目标子树的重心(要求除去 x 点时，它的 maxs 值最小，那么 x 就是这棵子树的重心了)，其中 tot 是这棵子树的总大小(节点个数)

    int maxs=tot-siz[x];//这棵子树中x 父亲那一支先赋给 maxs
    for (int i=head[x]; i; i=e[i].next)
    {
        int v=e[i].v;
        if (v!=fa&&!f[v])
        {
            maxs=max(maxs,siz[v]);
            dfs_root(v,tot,x);
        }
    }
    if (maxs<mins)
    {
        mins=maxs;
        root=x;
    }
}

int work(int x,int d)
{
    //返回以 x 为根的子树内长度小于等于 K 的路径数(两个端点都在子树内)
    //其实 d 在这里用处只有一个，是在做减法时方便把重心的儿子节点的 dis 先弄好，你也可以在分治的时候弄，不过就稍微有点麻烦了
    cnt=0;
    get_dis(x,d,0);
    sort(dis+1,dis+cnt+1);
    int daan=0,i=1,j=cnt;
    while (i<j)
    {
        while (i<j&&dis[i]+dis[j]>k)
        {
            j--;
        }
        daan+=j-i; //相当于选一条路径 i，另一条可以为 [i+1,j] 里任意一条路径，这样得到的两个点之间长度(经过重心的那条路径)肯定是小于等于 K 的
        i++;
    }
    return daan;
}

void dfs(int x)  //以 x 为重心分治一下
{
    cnt=0;
    mins=inf;
    get_size(x,0);
    dfs_root(x,siz[x],0);
    ans+=work(root,0);
    f[root]=1;
    for (int i=head[root]; i; i=e[i].next)//注意这里是以重心开始
    {
        int v=e[i].v;
        if (!f[v])
        {
            ans-=work(v,e[i].w);   //注意，这里 dis[o] 要先赋成 W[h](即它到重心的距离)
            dfs(v);
        }
    }
}

int main()
{
    while (~scanf("%d%d",&n,&k))
    {
        if (n==0&&k==0)
        {
            break;
        }
        for (int i=1,u,v,w; i<n; i++)
        {
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        dfs(1);
        printf("%d
",ans);
        num=ans=0;
        for (int i=1; i<=n; i++)
        {
            head[i]=f[i]=dis[i]=siz[i]=0;
        }
    }
    return 0;
}