【poj1741】 Tree

http://poj.org/problem?id=1741 (题目链接)

题意

　　给出一个n个节点的带权树，求树上距离不超过K的所有点对的个数。

solution 　

　　点分治裸题。所谓的点分治，就是对于一条路径，只有经过该点和不经过改点两种情况，所以我们可以通过找到树的重心，删去这个点，使树分成几棵小树，再递归处理。不经过的情况很好处理，直接递归到子树就可以了，关键是如何考虑经过的情况。

　　对于这道题，我们可以处理处所有点到当前子树重心的距离deep[]，然后对于当前子树的两个节点i，j只要满足deep[i]+deep[j]<=K，便符合条件，此时我们发现，这样算出的解是经过当前子树重心的情况，而当i，j都是重心的同一棵儿子节点子树上的点时，会重复计算，因为待会又会递归处理这一子树。所以我们还要减去这一部分的点对，方法同上，只是缩小了树的范围。接着递归处理出解。

　　代码hzwer那里模来的，很多细节感觉处理的很优秀，我加了点注释。

代码

// poj1741
#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdlib>
#include<cstdio>
#include<cmath>
#define MOD 1000000007
#define inf 2147483640
#define LL long long
#define free(a) freopen(a".in","r",stdin);freopen(a".out","w",stdout);
using namespace std;
inline int getint() {
    int x=0,f=1;char ch=getchar();
    while (ch>'9' || ch<'0') {if (ch=='-') f=-1;ch=getchar();}
    while (ch>='0' && ch<='9') {x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

const int maxn=10010;
struct edge {int to,w,next;}e[maxn<<2];
int head[maxn],vis[maxn],son[maxn],deep[maxn],f[maxn],d[maxn],n,cnt,root,sum,K,ans;

void insert(int u,int v,int w) {
    e[++cnt].to=v;e[cnt].next=head[u];head[u]=cnt;e[cnt].w=w;
    e[++cnt].to=u;e[cnt].next=head[v];head[v]=cnt;e[cnt].w=w;
}
void init() {
    cnt=ans=root=sum=0;
    memset(head,0,sizeof(head));
    memset(vis,0,sizeof(vis));
    memset(deep,0,sizeof(deep));
    for (int i=1;i<n;i++) {
        int u=getint(),v=getint(),w=getint();
        insert(u,v,w);
    }
}
void calroot(int u,int fa) {   //找重心
    f[u]=0;son[u]=1;
    for (int i=head[u];i;i=e[i].next) {
        if (fa==e[i].to || vis[e[i].to]) continue;
        calroot(e[i].to,u);
        son[u]+=son[e[i].to];
        f[u]=max(f[u],son[e[i].to]);
    }
    f[u]=max(f[u],sum-son[u]);
    if (f[u]<f[root]) root=u;
}
void caldeep(int u,int fa) {   //统计当前树中每个节点到重心的距离
    deep[++deep[0]]=d[u];
    for (int i=head[u];i;i=e[i].next) {
        if (e[i].to==fa || vis[e[i].to]) continue;
        d[e[i].to]=d[u]+e[i].w;
        caldeep(e[i].to,u);
    }
}
int cal(int u,int now) {
    d[u]=now;deep[0]=0;
    caldeep(u,0);
    sort(deep+1,deep+deep[0]+1);
    int t=0;
    for (int l=1,r=deep[0];l<r;) {   //这里统计答案的方法很优秀
        if (deep[l]+deep[r]<=K) t+=r-l,l++;
        else r--;
    }
    return t;
}
void work(int u) {
    ans+=cal(u,0);   //统计当前树中所有符合条件的点对
    vis[u]=1;
    for (int i=head[u];i;i=e[i].next) if (!vis[e[i].to]) {
            ans-=cal(e[i].to,e[i].w);   //减掉在同一棵子树中的符合条件的节点
            sum=son[e[i].to];
            root=0;
            calroot(e[i].to,0);
            work(root);   //递归分治处理子树
        }
}
int main() {
    while (scanf("%d%d",&n,&K)!=EOF && n) {
        init();
        sum=n;f[0]=inf;
        calroot(1,0);
        work(root);
        printf("%d\n",ans);
    }
    return 0;
}