《算法竞赛进阶指南》0x45点分治

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

给定一棵树，问树中长度小于等于K的点对的数量，经典的点分治问题，将这样的点对看做两种——经过根节点和不经过根节点，经过根结点的，可以通过子树划分，记录所有节点，然后按照指针方式进行搜索，可以卡过，时间复杂度大约是O(nlog^2n)。

代码：

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<string.h>
using namespace std;
const int maxn = 10010;
bool v[maxn],w[maxn];
int d[maxn],b[maxn],a[maxn],tot,cnt[maxn];
int head[maxn],nxt[maxn<<1],len[maxn<<1],ver[maxn<<1],t;//存图以及边的数量 
int n,k,s[maxn],Ans;
int ans,pos;//记录树的重心以及最大的子树的大小 
void addedge(int u,int v,int w){
    ver[++t]=v;
    len[t]=w;
    nxt[t]=head[u];
    head[u]=t;
}
void dfs_find(int S,int x){
    v[x]=1;
    s[x]=1;
    int max_part=0;
    for(int i=head[x];i;i=nxt[i]){
        int y=ver[i];
        if(v[y] || w[y])continue;//访问过或者是被删除
        dfs_find(S,y);
        s[x]+=s[y];
        max_part=max(max_part,s[y]);//找到删除x后最大的子树 
    }
    max_part=max(max_part,S-s[x]);
    if(max_part<ans){
        ans=max_part;
        pos=x;//重心 
    }
}
void dfs(int x){
    v[x]=1;
    for(int i=head[x];i;i=nxt[i]){
        int y=ver[i];
        int z=len[i];
        if(v[y] || w[y])continue;
        b[y]=b[x];
        a[++tot]=y;
        ++cnt[b[y]];
        d[y]=d[x]+z;
        dfs(y);
    }
}
bool cmp(int i,int j){//按照节点距离根节点的路径长度进行排序 
    return d[i]<d[j];
}
void work(int S,int x){
    memset(v,0,sizeof(v));
    ans=S;
    dfs_find(S,x);//找到重心 
    memset(d,0,sizeof(d));
    memset(cnt,0,sizeof(cnt));
    memset(v,0,sizeof(v));
    tot=1;//处理根节点 
    a[1]=pos;
    b[pos]=pos;
    w[pos]=1;//根节点被删除 
    ++cnt[pos];

//    访问pos的子树 
    for(int i=head[pos];i;i=nxt[i]){
        int y=ver[i];
        int z=len[i];
        if(v[y] || w[y])continue;
        a[++tot]=b[y]=y;
        ++cnt[b[y]];
        d[y]=z;//设置每个点距离根节点的距离
        dfs(y); //对每棵子树进行访问 
    }
    sort(a+1,a+tot+1,cmp);
    int l=1,r=tot;
    --cnt[b[a[1]]];//初始时刻维护的是[2,tot]区间上的信息
    while(l<r){
        while(d[a[l]]+d[a[r]]>k){
            cnt[b[a[r]]]--;
            r--;
        } 
        Ans+=r-l-cnt[b[a[l]]];//将同一棵子树上的r端减去 
        --cnt[b[a[l+1]]];//维护的区间发生变化 
        l++; 
    }
    int now=pos;//暂存这个重心，因为pos在处理子树的时候会变化 
    for(int i=head[now];i;i=nxt[i]){
        int y=ver[i];
        if(w[y])continue;
        work(s[y],y);//对每一个子树递归求解Ans 
    }
}
void solve(){
    t=0;
    memset(head,0,sizeof(head));
    for(int i=1;i<n;i++){
        int x,y,z;
        scanf("%d %d %d",&x,&y,&z);
        addedge(x,y,z);
        addedge(y,x,z);
    }
    memset(w,0,sizeof(w));
    Ans=0;
    work(n,1);//从任何一个点开始，每次都找重心进行操作
    cout<<Ans<<endl; 
}
int main(){
    while(cin>>n>>k && n && k)solve();
    return 0;
}