BZOJ 1468 Tree

点分治模板题。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define maxv 40050
#define maxe 80050
#define inf 1000000007
using namespace std;
struct edge
{
    int v,w,nxt;
}e[maxe];
int n,x,y,z,g[maxv],nume=0,k;
int f[maxv],root=0,size[maxv],dis[maxv],deep[maxv],ans=0,sum;
bool vis[maxv];
void addedge(int u,int v,int w)
{
    e[++nume].v=v;
    e[nume].w=w;
    e[nume].nxt=g[u];
    g[u]=nume;
}
void getroot(int now,int father)
{
    f[now]=0;size[now]=1;
    for (int i=g[now];i;i=e[i].nxt)
    {
        int v=e[i].v;
        if ((v!=father) && (vis[v]==false))
        {
            getroot(v,now);
            size[now]+=size[v];
            f[now]=max(f[now],size[v]);
        }
    }
    f[now]=max(f[now],sum-size[now]);
    if (f[now]<f[root]) root=now;
}
void getdis(int now,int father)
{
    deep[++deep[0]]=dis[now];
    for (int i=g[now];i;i=e[i].nxt)
    {
        int v=e[i].v;
        if ((vis[v]==false) && (v!=father))
        {
            dis[v]=dis[now]+e[i].w;
            getdis(v,now);
        }
    }
}
int calc(int x,int now)
{
    int cnt=0;
    deep[0]=0;dis[x]=now;
    getdis(x,0);
    sort(deep+1,deep+deep[0]+1);
    int l=1,r=deep[0];
    while (l<r)
    {
        if (deep[l]+deep[r]<=k)
        {
            cnt=cnt+r-l;
            l++;
        }
        else r--;
    }
    return cnt;
}
void work(int now)
{
    ans+=calc(now,0);
    vis[now]=true;
    for (int i=g[now];i;i=e[i].nxt)
    {
        int v=e[i].v;
        if (vis[v]==false)
        {
            ans-=calc(v,e[i].w);
            root=0;sum=size[v];
            getroot(v,root);
            work(root);
        }
    }
}
int main()
{
    memset(vis,false,sizeof(vis));
    scanf("%d",&n);
    for (int i=1;i<=n-1;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        addedge(x,y,z);
        addedge(y,x,z);
    }
    scanf("%d",&k);
    f[0]=inf;sum=n;
    getroot(1,0);
    work(root);
    printf("%d\n",ans);
    return 0;
}