POJ1741 Tree 树分治模板

http://poj.org/problem?id=1741

 

题意：一棵n个点的树，每条边有距离v，求该树中距离小于等于k的点的对数。

 

dis[y]表示点y到根x的距离,v代表根到子树根的距离；

那么不在同一棵子树中的两点i、j之间的距离为dis[i]+dis[j]； ① 

 

设得到这个距离的时间复杂度为O(w)；

如果我们层层如此递归即可得到所有的点对数量，可以证明复杂度为O(logn*w);

 

因为n的范围为(n<=10000)所以我们需要w与n近似；

 

那么此时问题转化为了如何在大约为O(n)的复杂度内得到①；

一个个计算不在同一子树中显然是麻烦的，如果选择先计算整棵树的点对数量然后去掉重复计数的点对数问题就可以得到简化；

 

如果只是计算一棵树下符合条件dis[j]+dis[i]<=k的点对数量，我们将距离sort，很容易在 log(树的大小) 的复杂度下把问题解决，再用几乎同样的时间减去每一棵子树中符合dis[j]+dis[i]+2*v<=k的点对的数量就可以得到答案。

那么总时间复杂度为O(n*logn*logn);

 

当然这只是理想情况，如果这棵树退化为一条链，复杂度则会变为O(n*n*logn)显然是超时的；

所以在每次递归前O(n)的复杂度找一下树的重心，

 

代码如下

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<cmath>
//using namespace std;
const int maxn=10010;
const double eps=1e-8;
const int modn=45989;
int n,k;
struct nod{
    int y,next;
    int v;
}e[maxn*2];
int head[maxn]={},siz[maxn]={},ma[maxn]={},dis[maxn]={};
int tot=0,tot1,root=0,now,ans;
bool vis[maxn]={};
void init(int x,int y,int v){
    e[++tot].y=y;
    e[tot].next=head[x];
    e[tot].v=v;
    head[x]=tot;
}
void getsiz(int x,int fa){ //某子树的大小
    int y;
    siz[x]=1;
    ma[x]=0;
    for(int i=head[x];i;i=e[i].next){
        y=e[i].y;
        if(y!=fa&&!vis[y]){
            getsiz(y,x);
            siz[x]+=siz[y];
            ma[x]=std::max(siz[y],ma[x]);
        }
    }
}
void cen(int r,int x,int fa){ //找重心
    int y;
    if(siz[r]-siz[x]>ma[x]){
        ma[x]=siz[r]-siz[x];
    }
    if(ma[x]<now){
        root=x;now=ma[x];
    }
    for(int i=head[x];i;i=e[i].next){
        y=e[i].y;
        if(y!=fa&&!vis[y]){
            cen(r,y,x);
        }
    }
}
void getdis(int x,int fa,int di){ //子树中各个点到该子树根的距离
    int v,y;
    dis[++tot1]=di;
    for(int i=head[x];i;i=e[i].next){
        v=e[i].v;y=e[i].y;
        if(y!=fa&&!vis[y]){
            getdis(y,x,di+v);
        }
    }
}
int sum(int x,int d){ //点对数
    tot1=0;
    getdis(x,0,d);
    std::sort(dis+1,dis+1+tot1);
    int i=1,j=tot1;
    int cnt=0;
    while(i<j){
        while(dis[i]+dis[j]>k&&i<j){
            j--;
        }
        cnt+=j-i;
        i++;
    }
    return cnt;
}
void dfs(int x){
    int y;
    now=maxn;root=x;
    getsiz(x,0);
    cen(x,x,0);
    ans+=sum(root,0);
    vis[root]=1;
    for(int i=head[root];i;i=e[i].next){
        y=e[i].y;
        if(!vis[y]){
            ans-=sum(y,e[i].v);
            dfs(y);
        }
    }
}
void yu(){ 
    memset(head,0,sizeof(head));
    memset(vis,0,sizeof(vis));
    tot=0;ans=0;
}
int main(){
    while((~scanf("%d%d",&n,&k))&&(n!=0||k!=0)){
        yu();
        int x,y,v;
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&x,&y,&v);
            init(x,y,v);
            init(y,x,v);
        }
        dfs(1);
        printf("%d
",ans);
    }
    return 0;
}