1468: Tree 点分治

Description
给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K

题解：

经典点分治……以前的模板好像有一点问题，这次加了个gettot函数，相当于多了个常数。

代码：

#include<bits/stdc++.h>
using namespace std;
#define LL long long
#define pa pair<int,int>
const int Maxn=40010;
const int inf=2147483647;
int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9')x=(x<<3)+(x<<1)+(ch^48),ch=getchar();
    return x*f;
}
int n,k,ans=0;
struct Edge{int y,d,next;}e[Maxn<<1];
int last[Maxn],len=0;
void ins(int x,int y,int d)
{
    int t=++len;
    e[t].y=y;e[t].d=d;
    e[t].next=last[x];last[x]=t;
}
bool vis[Maxn];
int smx,root,tot;
int gettot(int x,int fa)
{
    int re=0;
    for(int i=last[x];i;i=e[i].next)
    {
        int y=e[i].y;
        if(y==fa||vis[y])continue;
        re+=gettot(y,x);
    }return re;
}
int getroot(int x,int fa)
{
    int mx=0,size=0,rt;
    for(int i=last[x];i;i=e[i].next)
    {
        int y=e[i].y;
        if(y==fa||vis[y])continue;
        int t=getroot(y,x);
        mx=max(mx,t);size+=t;
    }
    mx=max(mx,tot-size-1);
    if(mx<smx)smx=mx,root=x;
    return size;
}
int la,a[Maxn];
void dfs(int x,int fa,int d)
{
    if(d<=k)a[++la]=d;
    for(int i=last[x];i;i=e[i].next)
    {
        int y=e[i].y;
        if(y==fa||vis[y])continue;
        dfs(y,x,d+e[i].d);
    }
}
int calc(int x,int d)
{
    la=0;dfs(x,0,d);
    sort(a+1,a+1+la);a[0]=0;
    int l=1,r=la,re=0;
    while(l<r)
    {
        if(a[l]+a[r]<=k)re+=(r-l),l++;
        else r--;
    }return re;
}
void solve(int x)
{
    vis[x]=true;
    ans+=calc(x,0);
    for(int i=last[x];i;i=e[i].next)
    {
        int y=e[i].y;
        if(vis[y])continue;
        ans-=calc(y,e[i].d);
        smx=inf,tot=gettot(y,x);getroot(y,x);
        solve(y);
    }
}
int main()
{
    n=read();
    for(int i=1;i<n;i++)
    {
        int x=read(),y=read(),d=read();
        ins(x,y,d),ins(y,x,d);
    }
    k=read();
    smx=inf,tot=n;getroot(1,0);solve(root);
    printf("%d",ans);
}