【acwing 252】树（树分治）

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/Coldfresh/article/details/88558805

给定一个有N个点（编号0,1,…,N-1）的树，每条边都有一个权值（不超过1000）。

树上两个节点x与y之间的路径长度就是路径上各条边的权值之和。

求长度不超过K的路径有多少条。

输入格式
输入包含多组测试用例。

每组测试用例的第一行包含两个整数N和K。

接下来N-1行，每行包含三个整数u,v,l，表示节点u与v之间存在一条边，且边的权值为l。

当输入用例N=0，K=0时，表示输入终止，且该用例无需处理。

输出格式
每个测试用例输出一个结果。

每个结果占一行。

数据范围
N≤10000
输入样例：
5 4
0 1 3
0 2 1
0 3 2
2 4 1
0 0
输出样例：
8

树分治喽，加个新标签
代码：

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<vector>
#include<queue>
#include<set>
#define maxn 10005
#define ll long long
#define maxx 1000005
#define mod 1000000007
#define INF 0x7f7f7f7f
using namespace std;
int n,k;
int head[maxn],_next[maxn<<1],to[maxn<<1],_w[maxn<<1];
int edge;
void addEdge(int u,int v,int w)
{
    to[++edge]=v,_w[edge]=w,_next[edge]=head[u],head[u]=edge;
    to[++edge]=u,_w[edge]=w,_next[edge]=head[v],head[v]=edge;
}
int vis[maxn];
int _size[maxn];
int core;
int num;
int _n;
void dfs0(int u,int fa)
{
    _size[u]=1;
    for(int i=head[u];i;i=_next[i])
    {
        int v=to[i];
        if(v==fa||vis[v])continue;
        dfs0(v,u);
        _size[u]+=_size[v];
    }
}
void dfs1(int u,int fa)
{
    int _max=0;
    for(int i=head[u];i;i=_next[i])
    {
        int v=to[i];
        if(v==fa||vis[v])continue;
        dfs1(v,u);
        _max=max(_max,_size[v]);
    }
    _max=max(_max,_n-_size[u]);
    if(_max<num)
    {
        num=_max;
        core=u;
    }
}
inline int getCore(int u,int fa)
{
    dfs0(u,fa);
    _n=_size[u];
    num=INF;
    core;
    dfs1(u,fa);
    //cout<<"core: "<<core<<endl;
    return core;
}
int d[maxn];
int b[maxn];
int cnt[maxn];
int a[maxn],tot;
int res;
void dfs2(int u,int fa)
{
    for(int i=head[u];i;i=_next[i])
    {
        int v=to[i];
        if(v==fa||vis[v])continue;
        d[v]=d[u]+_w[i];
        if(d[v]<=k)
        {
            a[++tot]=v;
            if(b[u]==0)b[v]=v;
            else b[v]=b[u];
            dfs2(v,u);
        }
    }
}
bool cmp(int x1,int x2)
{
    return d[x1]<d[x2];
}
int _count=0;
void calc(int p,int fa)
{
    d[p]=0,b[p]=0,tot=0;
    dfs2(p,fa);
    //cout<<"tot: "<<tot<<endl;
    res+=tot;
    if(tot<=1)return ;
    for(int i=1;i<=tot;i++)
        ++cnt[b[a[i]]];
    sort(a+1,a+1+tot,cmp);
    int L=1,R=tot;
    while(L<=R)
    {
        while(L<R&&d[a[L]]+d[a[R]]>k)
        {
            cnt[b[a[R]]]--;
            R--;
        }
        if(L<=R)
        {
            cnt[b[a[L]]]--;
            res+=R-L-cnt[b[a[L]]];
        }
        //cout<<cnt[b[a[L]]]<<"res: "<<res<<endl;
        L++;
    }
}
void solve(int p,int fa)
{
    calc(p,0);
    vis[p]=1;
    for(int i=head[p];i;i=_next[i])
    {
        int v=to[i];
        if(vis[v])continue;
        int _p=getCore(v,p);
        solve(_p,0);
    }
}
void init()
{
    memset(head,0,sizeof(head));
    edge=0;
    res=0;
    memset(vis,0,sizeof(vis));
}
int main()
{
    while(scanf("%d%d",&n,&k)==2)
    {
        if(n==0&&k==0)break;
        int x,y,w;
        init();
        for(int i=1;i<n;i++)
        {
            scanf("%d%d%d",&x,&y,&w);
            x++,y++;
            addEdge(x,y,w);
        }
        int p=getCore(1,0);
        solve(p,0);
        cout<<res<<endl;
    }
    return 0;
}