【树的点分治】【平衡树】[POJ1741]Tree

题目大意

给出一棵树，边带权，问有多少条长度路径<=k的路径。

分析

点分治

暴力做法为O(n2)我们可以将统计所有路径转化为递归查找通过某个特殊点的路径，从而降低复杂度。将这个点定为重心，复杂度可以降低为O(nlogn)

复杂度分析

重心有一个性质，如果以它为根，它的最大一个子树的大小不超过2tot。如果超过，那么在那棵子树里面一定有一个点它最大的子树小于这个点的最大子树，那么我们现在所选择的点就不是重心。那么我们经过logn次分治后得到空树，每一层的复杂度为O(n)，总的复杂度为O(nlogn)

这道题

做法

我们每次对重心的所有子树进行遍历，求出重心到这些点的距离，用一棵平衡树维护。遍历完一棵子树后将这棵子树中所有距离一次性加入平衡树中。在遍历一棵子树时，设当前点到重心的距离为dep,就可以很快地求出之前遍历的子树中有多少点到重心的距离≤k−dep，复杂度再加一个logn，总的复杂度O(nlog2n)
然后切断重心和子树的边，对子树进行递归处理。

正确性

对于通过重心的路径，我们显然已经计算到了。
没有通过的，显然只会存在于一棵子树中，会在递归处理这棵子树时被处理到。

代码

#include<cstdio>
#include<ctime>
#include<cstring>
#include<algorithm>
using namespace std;
#define MAXN 10000
void Read(int &x){
    char c;
    while(c=getchar(),c!=EOF)
        if(c>='0'&&c<='9'){
            x=c-'0';
            while(c=getchar(),c>='0'&&c<='9')
                x=x*10+c-'0';
            ungetc(c,stdin);
            return;
        }
    exit(0);
}
struct node{
    int v,wt;
    node *next;
}*adj[MAXN+10],edge[MAXN*2+10],*ecnt=edge;
int n,k,size[MAXN+10],mxs[MAXN+10],heavy,tmp[MAXN+10],ans,cnt;
bool vis[MAXN+10];
void addedge(int u,int v,int wt){
    node *p=++ecnt;
    p->v=v;
    p->wt=wt;
    p->next=adj[u];
    adj[u]=p;
}
void read(){
    int i,u,v,wt;
    for(i=1;i<n;i++){
        Read(u),Read(v),Read(wt);
        addedge(u,v,wt);
        addedge(v,u,wt);
    }
}
void dfs(int u,int fa,int tot){
    size[u]=1,mxs[u]=0;
    for(node *p=adj[u];p;p=p->next)
        if(p->v!=fa&&!vis[p->v]){
            dfs(p->v,u,tot);
            size[u]+=size[p->v];
            mxs[u]=max(mxs[u],size[p->v]);
        }
    mxs[u]=max(mxs[u],tot-size[u]);
    if(mxs[u]<mxs[heavy])
        heavy=u;
}
struct Treap_node{
    int val,pri,cnt,size;
    Treap_node *ch[2];
}tree[MAXN+10],*tcnt=tree,*root=0;
inline int Get_size(Treap_node *p){
    return p?p->size:0;
}
void update(Treap_node *p){
    p->size=Get_size(p->ch[0])+Get_size(p->ch[1])+p->cnt;
    return;
}
void Rotate(Treap_node *&x,bool d){
    Treap_node *y=x->ch[!d];
    x->ch[!d]=y->ch[d];
    y->ch[d]=x;
    update(x);
    x=y;
}
void insert(Treap_node *&p,int val){
    if(!p){
        p=++tcnt;
        p->val=val;
        p->pri=rand();
        p->cnt=p->size=1;
        p->ch[0]=p->ch[1]=0;
        return;
    }
    if(val==p->val){
        p->cnt++;
        p->size++;
        return;
    }
    bool d=val>p->val;
    insert(p->ch[d],val);
    if(p->ch[d]->pri>p->pri)
        Rotate(p,!d);
    update(p);
}
int find_size(Treap_node *p,int val){
    if(!p)
        return 0;
    if(val==p->val)
        return Get_size(p->ch[0])+p->cnt;
    if(val>p->val)
        return Get_size(p->ch[0])+p->cnt+find_size(p->ch[1],val);
    return find_size(p->ch[0],val);
}
void dfs2(int u,int fa,int dep){
    size[u]=1;
    ans+=find_size(root,k-dep);
    tmp[++cnt]=dep;
    for(node *p=adj[u];p;p=p->next)
        if(p->v!=fa&&!vis[p->v]){
            dfs2(p->v,u,dep+p->wt);
            size[u]+=size[p->v];
        }
}
void solve(int u){
    heavy=0;
    dfs(u,0,size[u]);
    root=0,tcnt=tree;
    vis[heavy]=1;
    int i;
    insert(root,0);
    for(node *p=adj[heavy];p;p=p->next){
        cnt=0;
        if(!vis[p->v]){
            dfs2(p->v,heavy,p->wt);
            for(i=1;i<=cnt;i++)
                insert(root,tmp[i]);
        }
    }
    for(node *p=adj[heavy];p;p=p->next)
        if(!vis[p->v])
            solve(p->v);
}
inline void init(){
    memset(adj,0,sizeof adj);
    ecnt=edge;
    ans=0;
    memset(vis,0,sizeof vis);
}
int main()
{
    srand(254587867);
    while(Read(n),Read(k),n&&k){
        init();
        read();
        size[1]=n;
        mxs[0]=0x7fffffff;
        solve(1);
        printf("%d\n",ans);
    }
}