POJ1741 Tree（BZOJ1468）

点分治

POJ题目传送门

BZOJ题目传送门

经典点分治题。

每次求出重心，计算满足条件的点对（距离≤k），再减去多余的点对（不经过当前根的点对），然后递归处理其子树，累加即为答案。

代码：

#include<cstdio>
#include<cstring>
#include<algorithm>
#define MAXN 10000
using namespace std;
struct edge{
    int next,to,dis;
}ed[MAXN*2+5];
int n,m,k,rt,ans,t,p;
int h[MAXN+5],sum[MAXN+5],dis[MAXN+5],ma[MAXN+5],dep[MAXN+5];
bool f[MAXN+5];
void addedge(int x,int y,int z){
    ed[++k].next=h[x]; ed[k].to=y; ed[k].dis=z; h[x]=k;
}
void dfsrt(int x,int fa){//找重心
    sum[x]=1; ma[x]=0;
    for (int i=h[x];i;i=ed[i].next)
        if (ed[i].to!=fa&&!f[ed[i].to]){
            int v=ed[i].to; dfsrt(v,x);
            sum[x]+=sum[v]; ma[x]=max(ma[x],sum[v]);
        }
    ma[x]=max(ma[x],t-sum[x]);
    if (ma[x]<ma[rt]) rt=x;
}
void dfsdep(int x,int fa){//算距离
    dep[++p]=dis[x];
    for (int i=h[x];i;i=ed[i].next)
        if (ed[i].to!=fa&&!f[ed[i].to]){
            dis[ed[i].to]=dis[x]+ed[i].dis;
            dfsdep(ed[i].to,x);
        }
}
int dfssum(int x,int v){//求满足条件点对数
    dis[x]=v; p=0; dfsdep(x,0);
    sort(dep+1,dep+p+1);
    int l=1,r=p,ret=0;
    while (l<r)
        if (dep[l]+dep[r]<=m) ret+=r-l,l++; 
        else r--;
    return ret; 
}
void dfsans(int x){//答案
    f[x]=true; ans+=dfssum(x,0);//总方案
    for (int i=h[x];i;i=ed[i].next)
        if (!f[ed[i].to]){
            int v=ed[i].to;
            ans-=dfssum(v,ed[i].dis);//多余方案
            rt=0; t=sum[v];
            dfsrt(v,0); dfsans(rt);
        }
}
int main(){
    while(scanf("%d%d",&n,&m)&&(n||m)){
        memset(f,false,sizeof(f));
        memset(h,0,sizeof(h)); k=ans=0; 
        for (int i=1;i<n;i++){
            int u,v,d; scanf("%d%d%d",&u,&v,&d);
            addedge(u,v,d); addedge(v,u,d);
        }
        ma[(rt=0)]=0x7fffffff; t=n;
        dfsrt(1,0); 
        dfsans(rt);
        printf("%d\n",ans);
    }
    return 0;
}