Cogs 1714. [POJ1741][男人八题]树上的点对(点分治)

1. [POJ1741][男人八题]树上的点对
   ★★★ 输入文件：poj1741_tree.in 输出文件：poj1741_tree.out 简单对比
   时间限制：1 s 内存限制：256 MB
   【题目描述】
   给一棵有n个节点的树，每条边都有一个长度（小于1001的正整数）。
   定义dist(u,v)=节点u到节点v的最短路距离。
   给出一个整数k，我们称顶点对(u,v)是合法的当且仅当dist(u,v)不大于k。
   写一个程序，对于给定的树，计算有多少对顶点对是合法的。
   【输入格式】
   输入包含多组数据。
   每组数据的第一行有两个整数N，K（N<=10000）。接下来N-1行每行有三个整数u,v,l，代表节点u和v之间有一条长度l的无向边。
   输入结束标志为N=K=0.
   【输出格式】
   对每组数据输出一行一个正整数，即合法顶点对的数量。
   【样例输入】
   5 4
   1 2 3
   1 3 1
   1 4 2
   3 5 1
   0 0
   【样例输出】
   8
   【来源】
   **POJ 1741 Tree
   男人八题 Problem E**

/*
乱搞80分2333.
T了.
处理出每个树的贡献,再减去在每颗子树中的贡献.
可能是让链卡掉了QWQ.
每个点至多有logn个父亲.
也就是在logn个子树中.
复杂度应该是O(Nlogn2)的.
如果是链的话就变成了O(N^2).
*/
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#define MAXN 10001
using namespace std;
int n,k,tot,ans,cut,dis[MAXN],maxsize,s[MAXN],a[MAXN],pos[MAXN],end[MAXN],fa[MAXN],head[MAXN],deep[MAXN],size[MAXN];
struct edge{int v,next,x;}e[MAXN*2];
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*10+ch-48,ch=getchar();
    return x*f;
}
void add(int u,int v,int x)
{
    e[++cut].v=v;e[cut].x=x;e[cut].next=head[u];head[u]=cut;
}
void dfs(int u)
{
    s[++maxsize]=dis[u];pos[u]=maxsize;
    for(int i=head[u];i;i=e[i].next)
    {
        if(!fa[e[i].v])
          deep[e[i].v]=deep[u]+1,
          fa[e[i].v]=u,dis[e[i].v]=dis[u]+e[i].x,dfs(e[i].v);
    }
    end[u]=maxsize;
    return ;
}
int erfen(int l,int r)
{
    int total=0;
    while(l<r)
    {
        if(a[l]+a[r]<=k) total+=r-l,l++;
        else r--;
    }
    return total;
}
int get(int u,int d)
{
    int sum=0;tot=0;
    for(int i=pos[u];i<=end[u];i++) a[++tot]=s[i]-dis[u]+d;
    sort(a+1,a+tot+1);
    return erfen(1,tot);
}
void get_dis(int u)
{
    ans+=get(u,0);
    for(int i=head[u];i;i=e[i].next)
      if(fa[e[i].v]==u) ans-=get(e[i].v,e[i].x),get_dis(e[i].v);
    return ;
}
void slove()
{
    fa[1]=1;dfs(1);
    for(int i=1;i<=n;i++) a[i]=s[i];
    sort(a+1,a+n+1);
    ans+=erfen(1,n);
    for(int i=head[1];i;i=e[i].next)
      if(fa[e[i].v]==1) ans-=get(e[i].v,e[i].x),get_dis(e[i].v);
    return ;
}
void Clear()
{
    memset(head,0,sizeof head);
    memset(fa,0,sizeof fa);
    cut=maxsize=tot=ans=0;
}
int main()
{
    freopen("poj1741_tree.in","r",stdin);
    freopen("poj1741_tree.out","w",stdout);
    int x,y,z;
    while(true)
    {
        Clear();n=read(),k=read();
        if(!n&&!k) break;
        for(int i=1;i<=n-1;i++)
        {
            x=read(),y=read(),z=read();
            add(x,y,z),add(y,x,z);
        }
        slove();
        printf("%d\n",ans);
    }
    return 0;
}

/*
点分治. 
找树的重心.
第一次打,抄的黄学长的.
重心是使得与该点相邻的联通块点数
最大的最小化的那个点.
显然如果用一颗树的重心作根的话,
可以使子树最大的最小化. 
复杂度应该是O(Nlogn2)的.
*/
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#define MAXN 10001
using namespace std;
int n,k,tot,ans,cut,f[MAXN],sum,root,dis[MAXN],maxsize,a[MAXN],head[MAXN],deep[MAXN],size[MAXN];
bool b[MAXN];
struct edge{int v,next,x;}e[MAXN*2];
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*10+ch-48,ch=getchar();
    return x*f;
}
void add(int u,int v,int x)
{
    e[++cut].v=v;e[cut].x=x;e[cut].next=head[u];head[u]=cut;
}
int erfen(int l,int r)
{
    int total=0;
    while(l<r)
    {
        if(a[l]+a[r]<=k) total+=r-l,l++;
        else r--;
    }
    return total;
}
void get_root(int u,int fa)
{
    size[u]=1;f[u]=0;
    for(int i=head[u];i;i=e[i].next)
    {
        if(e[i].v==fa||b[e[i].v]) continue;
        get_root(e[i].v,u);
        size[u]+=size[e[i].v];
        f[u]=max(f[u],size[e[i].v]);
    }
    f[u]=max(f[u],sum-size[u]);//u上方的联通块. 
    if(f[u]<f[root]) root=u;
}
void get_deep(int u,int fa)
{
    a[++tot]=dis[u];
    for(int i=head[u];i;i=e[i].next)
    {
        if(e[i].v==fa||b[e[i].v]) continue;
        dis[e[i].v]=dis[u]+e[i].x;
        get_deep(e[i].v,u);
    }
    return ;
}
int get(int u,int d)
{
    dis[u]=d;tot=0;
    get_deep(u,0);
    sort(a+1,a+tot+1);
    return erfen(1,tot);
}
void slove(int u)
{
    b[u]=true;ans+=get(u,0);
    for(int i=head[u];i;i=e[i].next)
    {
        if(b[e[i].v]) continue;
        ans-=get(e[i].v,e[i].x);
        sum=size[e[i].v];
        root=0;
        get_root(e[i].v,root);
        slove(root);
    }
    return ;
}
void Clear()
{
    memset(head,0,sizeof head);
    memset(b,0,sizeof b);
    memset(f,0,sizeof f);
    root=cut=maxsize=tot=ans=0;
}
int main()
{
    freopen("poj1741_tree.in","r",stdin);
    freopen("poj1741_tree.out","w",stdout);
    int x,y,z;
    while(true)
    {
        Clear();n=read(),k=read();
        if(!n&&!k) break;
        for(int i=1;i<=n-1;i++)
        {
            x=read(),y=read(),z=read();
            add(x,y,z),add(y,x,z);
        }
        f[0]=1e9;
        sum=n,get_root(1,root);
        slove(root);
        printf("%d\n",ans);
    }
    return 0;
}