Tree POJ - 1741 (点分治)

板子

//#include<bits/stdc++.h>
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<cstring>
#include<queue>
#include<set>
#include<map>
#include<vector>

#define inf 0x3f3f3f3f
#define mem(a,b) memset(a,b,sizeof(a))
#define ll long long
#define sd(x) scanf("%d",&(x))
#define sl(x) scanf("%lld",&(x))
#define slf(x) scanf("%lf",&(x))
#define scs(s) scanf("%s",s)
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define per(i,a,b) for(int i=a;i>=b;i--)
#define lowbit(x) x&(-x)
#define ls now<<1
#define rs now<<1|1
#define lson l,mid,ls
#define rson mid+1,r,rs
#define All L,R

using namespace std;

const int maxn=2e6+10;

struct edge{
    int u,v,w,next;
}e[maxn];

int top,mub,root,Max,ans,n,k,cnt,tot=0;
int g[maxn],Size[maxn],dist[maxn];
bool vis[maxn];

void make(int u,int v,int w)
{
    e[++tot]=(edge){u,v,w,g[u]};
    g[u]=tot;
}

void init()
{
    mem(g,0);
    tot=0;
    top=0;
    ans=0;
    mem(vis,0);
}

void get_root(int u,int f)
{
    Size[u]=1;
    int ma=0;
    for(int i=g[u];i;i=e[i].next)
    {
        int v=e[i].v;
        if(v!=f&&!vis[v])
        {
            get_root(v,u);
            Size[u]+=Size[v];
            ma=max(ma,Size[v]);
        }
    }
    ma=max(ma,cnt-Size[u]);
    if(ma<Max)
    {
        Max=ma;
        root=u;
    }
}


void get_dis(int u,int f,int d)
{
    dist[mub++]=d;
    for(int i=g[u];i;i=e[i].next)
    {
        int v=e[i].v,w=e[i].w;
        if(!vis[v]&&v!=f)
        {
            get_dis(v,u,d+w);
        }
    }
}

int calc(int u,int w)
{
    int res=0;
    mub=0;
    get_dis(u,-1,w);
    sort(dist,dist+mub);
    int i=0,j=mub-1;
    while(i<j)
    {
        while(i<j&&dist[i]+dist[j]>k) j--;
        res+=j-i;
        i++;
    }
    return res;
}


void work(int u)
{
    Max=n;
    get_root(u,-1);
    cnt=Size[u];
    ans+=calc(root,0);
    vis[root]=true;
//    cout<<"root: "<<root<<endl;
    for(int i=g[root];i;i=e[i].next)
    {
        int v=e[i].v,w=e[i].w;
        if(!vis[v])
        {
            ans-=calc(v,w);
            work(v);
        }
    }
}


int main()
{

    while(~sd(n))
    {
        sd(k);
        if(!n&&!k) break;
        init();
        int u,v,w;
        rep(i,1,n-1)
        {
            sd(u),sd(v),sd(w);
            make(u,v,w);
            make(v,u,w);
        }
        work(1);
        cout<<ans<<endl;

    }
    return 0;
}