poj 1741 Tree（树的点分治）

poj 1741 Tree（树的点分治）

给出一个n个结点的树和一个整数k，问有多少个距离不超过k的点对。

首先对于一个树中的点对，要么经过根结点，要么不经过。所以我们可以把经过根节点的符合点对统计出来。接着对于每一个子树再次运算。如果不用点分治的技巧，时间复杂度可能退化成\(O(n^2)\)（链）。如果对于子树重新选根，找到树的重心，就一定可以保证时间复杂度在\(O(nlogn)\)内。

具体技巧是：首先选出树的重心，将重心视为根。接着计算出每个结点的深度，以此统计答案。由于子树中可能出现重复情况，需要在子树中相应的减去一部分ans。具体实现在代码中。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int maxn=1e4+5;

struct Graph{
    struct Edge{
        int to, next, v; Graph *bel;
        Edge& operator ++(){
            return *this=bel->edge[next]; }
    }edge[maxn*2];
    int cntedge, fir[maxn];
    void addedge(int x, int y, int v){
        Edge &e=edge[++cntedge];
        e.to=y; e.next=fir[x]; e.v=v;
        fir[x]=cntedge; e.bel=this;
    }
    Edge& getlink(int x){ return edge[fir[x]]; }
    void RESET(){ cntedge=0; memset(fir, 0, sizeof(fir)); }
}g;

int n, k, size[maxn], w[maxn], dep[maxn];
int cnt[maxn], tail;
bool done[maxn];

//获取子树大小
int getsize(int now, int par, int num){
    size[now]=0; w[now]=0;
    Graph::Edge e=g.getlink(now);
    for (; e.to; ++e){
        if (e.to==par||done[e.to]) continue;
        size[now]+=getsize(e.to, now, num);
        w[now]=max(w[now], size[e.to]);
    }
    ++size[now]; w[now]=max(w[now], num-size[now]);
    return size[now];
}

//获取根的位置
int getroot(int now, int par){
    Graph::Edge e=g.getlink(now);
    int root=now, tmp=now;
    for (; e.to; ++e){
        if (e.to==par||done[e.to]) continue;
        tmp=getroot(e.to, now);
        if (w[tmp]<w[root]) root=tmp; //mdzzle
    }
    return root;
}

//获取子树中结点深度，并创造深度数组
void getdep(int now, int par, int step){
    Graph::Edge e=g.getlink(now);
    for (; e.to; ++e)
        if (e.to!=par&&!done[e.to])
        getdep(e.to, now, step+e.v);
    dep[now]=step;
    cnt[tail++]=step;
}

//通过深度数组统计和小于等于k的数对
int getans(int k, int tail){
    sort(cnt, cnt+tail);
    --tail; int ans=0;
    for (int l=0; l<tail; ++l){
        while (cnt[l]+cnt[tail]>k&&l<tail) --tail;
        ans+=tail-l;
    }
    return ans;
}

int solve(int now, int num){
    getsize(now, 0, num); //获取当前树的子树大小
    if (size[now]==1) return 0;
    now=getroot(now, 0); //凭借子树大小找到重心
    tail=0; //把深度数组复原
    getdep(now, 0, 0); //获取每个结点的深度，构建深度数组
    //获取答案，别忘记减掉重复的点对
    int ans=getans(k, tail);
    done[now]=true; //堵住当前点
    Graph::Edge e=g.getlink(now);
    for (; e.to; ++e) if (!done[e.to]){
        tail=0; getdep(e.to, 0, 0);
        ans-=getans(k-e.v*2, tail);
        ans+=solve(e.to, size[e.to]);
    }
    return ans;
}

int main(){
    while (~scanf("%d%d", &n, &k)&&n&&k){
        int t1, t2, t3; g.RESET();
        for (int i=1; i<=n; ++i) done[i]=false;
        for (int i=1; i<n; ++i){
            scanf("%d%d%d", &t1, &t2, &t3);
            g.addedge(t1, t2, t3);
            g.addedge(t2, t1, t3);
        }
        printf("%d\n", solve(1, n));
    }
    return 0;
}