【poj1741】tree（点分治例题）

题目描述：给一棵n个节点的树，每条边上有一个距离。定义d(u,v)为u到v的最小距离。给定k值，求有多少点对（u，v）使u到v的距离小于等于k。

思路：点分治例题，对每个点算出他到其他各点距离dis，再对每个点算他能和其他点组成的对数，即该点对答案的贡献。

学习博客：https://blog.csdn.net/qq_39553725/article/details/77542223

代码：

#include<iostream>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<sstream>
#include<vector>
#include<stack>
#include<deque>
#include<cmath>
#include<map>
#include<queue>
#include<bitset>
//#include<hash_map>
#define sd(x) scanf("%d",&x)
#define lsd(x) scanf("%lld",&x)
#define ms(x,y) memset(x,y,sizeof x)
#define fu(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define all(a) a.begin(),a.end()
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
using namespace std;
//using namespace __gnu_cxx;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
const int maxn=1e4+79;
const int mod=1e9+7;
const int INF=1e9+7;
const double pi=acos(-1);
struct od
{
    int to,nxt,len;
}t[maxn*4];
/*
int su[maxn*4];
//--------------------分割线-------------------------
ll mp(ll x,ll n){ll res=1; while(n){if(n&1) res=res*x%mod;x=x*x%mod;n>>=1;} return res;}
void add(int i,int k){ while(i<=maxn) su[i]+=k,i+=i&-i;}
ll gsum(int i){ll res=0;while(i>0) res+=su[i],i-=i&-i; return res;}
*/

int root,mx,mxson[maxn],sim[maxn],simer=0,cnt;
int dis[maxn],head[maxn],ans,tot,n,k;
bool vis[maxn];
void init()
{
    cnt=ans=0;
    simer=n,mx=INF;
    ms(vis,0);ms(head,0);ms(sim,0);
}
void add(int u,int v,int c)
{
    ++cnt;
    t[cnt].to=v;
    t[cnt].len=c;
    t[cnt].nxt=head[u];
    head[u]=cnt;
}
void getroot(int u,int fa)
{
    //获得树的重心
    sim[u]=1;mxson[u]=0;
    for(int i=head[u];i;i=t[i].nxt)
    {
        int v=t[i].to;
        if(vis[v]||v==fa) continue;
        getroot(v,u);
        sim[u]+=sim[v];
        mxson[u]=max(mxson[u],sim[v]);
    }
    mxson[u]=max(mxson[u],simer-sim[u]);
    if(mxson[u]<mx) root=u,mx=mxson[u];
}
void getdis(int u,int fa,int dist)
{
    //得到每个点的dis
    dis[++tot]=dist;
    for(int i=head[u];i;i=t[i].nxt)
    {
        int v=t[i].to;
        if(vis[v]||v==fa) continue;
        getdis(v,u,dist+t[i].len);
    }
    return;
}
int cal(int u,int len)
{
    //得到每个点的贡献
    tot=0;ms(dis,0);//初始化
    getdis(u,0,len);
    sort(dis+1,dis+tot+1);
    int l=1,r=tot,res=0;
    while(l<=r)
    {
        //l和l+1~r都能配对
        if(dis[l]+dis[r]<=k) res+=r-l,l++;
        else r--;
    }
    return res;
}
void dive(int u)
{
    ans+=cal(u,0);
    vis[u]=1;
    for(int i=head[u];i;i=t[i].nxt)
    {
        int v=t[i].to;
        if(vis[v]) continue;
        ans-=cal(v,t[i].len);
        simer=sim[v];root=0;mx=INF;
        getroot(v,0);
        dive(root);
    }
    return;
}
int main()
{
    while(1)
    {
        sd(n),sd(k);
        if(n==k&&n==0) break;
        init();
        fu(i,1,n-1)
        {
            int u,v,len;sd(u),sd(v);sd(len);
            add(u,v,len);add(v,u,len);
        }
        getroot(1,0);
        dive(root);
        printf("%d\n",ans);
    }
    return 0;
}