poj1741（树的点分治）

题目连接：POJ - 1741

看了好长时间才明白了点......

网上讲解很多但感觉都不够详细。。。大概是太弱了吧-_-||

学通了再回来写详解。。。

#include<iostream>
#include<cstring>
#include<cstdio>
#include<algorithm>
#define LL long long
using namespace std;
const int N=1e5+10;
const int inf=1e9+10;
int n,k;
int rt,ans,sumnode,siz[N],vis[N],d[N];
int dep[N]; //路径长度//dep[0]为子节点个数
int f[N]; //重心
struct edge
{
    int v,w,nex;
}e[N*2];
int head[N];
int cnt;
void add(int u,int v,int w)
{
    e[cnt].v=v;
    e[cnt].w=w;
    e[cnt].nex=head[u];
    head[u]=cnt++;
}
void getrt(int u,int fa) //找重心
{
    siz[u]=1;
    f[u]=0;
    for(int i=head[u];i!=-1;i=e[i].nex)
    {
        int v=e[i].v;
        if(v==fa||vis[v]) continue;
        getrt(v,u);
        siz[u]+=siz[v];
        f[u]=max(f[u],siz[v]);
    }
    f[u]=max(sumnode-siz[u],f[u]);
    if(f[u]<f[rt]) rt=u;
}

void getdep(int u,int fa)  //获取子树所有节点与根的距离
{
    dep[++dep[0]]=d[u];
    for(int i=head[u];i!=-1;i=e[i].nex)
    {
        int v=e[i].v;
        if(v==fa||vis[v]) continue;
        d[v]=d[u]+e[i].w;
        getdep(v,u);
    }
}

int cal(int u,int now) //计算当前以重心u为根的子树下所有情况的答案
{
    d[u]=now;
    dep[0]=0;
    getdep(u,0);
    sort(dep+1,dep+dep[0]+1);
    int all=0;
    for(int l=1,r=dep[0];l<r;)
    {
        if(dep[l]+dep[r]<=k) {all+=r-l;l++;}
        else r--;
    }
    return all;
}

void work(int u) //以u为重心计算
{
    vis[u]=1;
    ans+=cal(u,0);
    for(int i=head[u];i!=-1;i=e[i].nex)
    {
        int v=e[i].v;
        if(vis[v]) continue;
        ans-=cal(v,e[i].w);
        sumnode=siz[v];
        rt=0;
        getrt(v,u);
        work(rt);
    }
}
int u,v,w;
int main()
{
    while(scanf("%d%d",&n,&k)!=EOF)
    {
        if(!n&&!k) break;
        memset(head,-1,sizeof(head));
        memset(vis,0,sizeof(vis));
        cnt=0;
        for(int i=0;i<n-1;i++)
        {
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        rt=ans=0;
        sumnode=n;
        f[0]=inf;
        getrt(1,0);
        work(rt);
        printf("%d\n",ans);
    }
}