AcWing 252. 树(点分治)

题目大意：给定n个点的树，求树上路径长度不大于k的路径条数

题解：点分治模板

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;

const int N=10010;

int n,m,js;

int sumedge,head[N];

int p[N],q[N];
//P存所有结点到重心的路径，q存当前子树所有结点到重心的路径
bool st[N];

struct Edge
{
    int x,y,z,nxt;
    Edge(int x=0,int y=0,int z=0,int nxt=0):
    x(x),y(y),z(z),nxt(nxt){}
}edge[N<<1];

void add(int x,int y,int z)
{
    edge[++sumedge]=Edge(x,y,z,head[x]);
    head[x]=sumedge;
}

int get_size(int now,int fa)//求now的子树大小
{
    if(st[now]) return 0;
    int res=1;
    for(int i=head[now];i;i=edge[i].nxt)
    {
        int v=edge[i].y;
        if(v==fa) continue;
        res+=get_size(v,now);
    }
    return res;
}

int get_wc(int now,int fa,int tot,int& wc) //求重心
{
    if(st[now]) return 0;
    int sum=1,ms=0;
    for(int i=head[now];i;i=edge[i].nxt)
    {
        int v=edge[i].y;
        if(v==fa) continue;
        int t=get_wc(v,now,tot,wc);
        sum+=t;
        ms=max(ms,t);
    }
    ms=max(ms,tot-sum);
    if(ms<=tot/2) wc=now;
    return sum;
}

void get_dist(int now,int fa,int dist,int& qt)
{
    if(st[now]) return;
    q[++qt]=dist;
    for(int i=head[now];i;i=edge[i].nxt)
    {
        int v=edge[i].y;
        if(v==fa) continue;
        get_dist(v,now,dist+edge[i].z,qt);
    }
}

int get(int a[],int k)
{
    sort(a+1,a+k+1);
    int res=0;
    for(int i=k,j=0;i>=1;i--)
    {
        while(j+1<i&&a[j+1]+a[i]<=m) j++;
        j=min(j,i-1);
        res+=j;
    }
    return res;
}

int calc(int now)//处理now所在的树
{
  //  js++;
  //  if(js>4*n) exit(0);
    if(st[now]) return 0;
    int res=0;
    get_wc(now,-1,get_size(now,-1),now);
    st[now]=true;//now为重心
    int pt=0;
    for(int i=head[now];i;i=edge[i].nxt)
    {
        int v=edge[i].y;
        int qt=0;
        get_dist(v,-1,edge[i].z,qt);//v所在子树到重心距离
        res-=get(q,qt);//减去子树中符合要求的部分
        for(int k=1;k<=qt;k++)//?
        {
            p[++pt]=q[k];
            if(q[k]<=m) res++;
        }
    }
    res+=get(p,pt);
    for(int i=head[now];i;i=edge[i].nxt)
    {
        int v=edge[i].y;
    //    cout<<now<<"-----"<<v<<endl;
        res=res+calc(v);//递归部分
    }
    return res;
}

int main()
{
    while(cin>>n>>m)
    {
        if(!n&&!m) break;
        memset(head,0,sizeof(head));
        memset(st,0,sizeof(st));
        sumedge=0;
        for(int i=1;i<n;i++)
        {
            int x,y,z;
            scanf("%d%d%d",&x,&y,&z);
            add(x,y,z);
            add(y,x,z);
        }
        printf("%d
",calc(0));
    }
    return 0;
}