POJ1741--Tree[平衡树]

题目描述

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.

题解

题目描述大概是这样的：
给出一棵树，给定一个K，求这棵树中的点对(u,v)的个数，满足u,v的距离小于等于K。
这题比较正常的写法是点分，但是由于本蒟蒻不会点分，所以这里介绍一种用平衡树来做的方法：
我们最想维护的是什么？
我们最想维护每一棵子树中的所有节点到子树根节点的距离，如果我们能够维护出这个信息，很容易可以求出到子树根的距离小于等于K的个数，但是，还有另外一种情况就是u,v分别在根节点的两棵子树中，这样的情况就需要枚举其中一棵子树中的节点，假设这个节点到根节点的距离为d，然后在另一棵子树中查找到根节点的距离小于等于K-d的节点个数。但是这个过程涉及到枚举两棵子树再枚举其中一棵子树中的所有节点，显然复杂度太高。
如何改进？
先要介绍一种用来合并两棵Treap的方法，启发式合并:
若节点总数n为常数，对于这样的合并：每次把节点数少的Treap中的元素一个一个加到节点数大的Treap中去，把所有Treap合并为一棵的复杂度为O(nlogn)。证明也很显然，对于每一个节点，它每参加一次合并，它所在的Treap的节点数就*2，那么最终的Treap的节点数为n，所以，它最多参加logn次合并，每一个节点都是这样，所以总共的复杂度为O(nlogn)。
有了启发式合并，我们就可以快速的实现两棵Treap的合并，所以，对于这题，可以利用每一个节点最多被合并logn次这个性质，放心大胆的给每一个节点都建立一个Treap来维护以它为根的子树到它的距离，并且合并的总复杂度也是O(nlogn)的，完全可以接受。累计答案时，可以直接将子树与根节点所在的Treap合并，边合并边在另一棵子树中累计答案。由于要先维护子树，所以需要自底向上来遍历树，可以利用DFS天然的自底向上的性质来实现。
代码如下：

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#define maxn 200006
#define maxe 20006
using namespace std;
struct node{
    node *ch[2];
    int key,f,s,cnt,tag;
    void maintain(){s=cnt+ch[0]->s+ch[1]->s;}
    void pushdown(){
        ch[0]->tag+=tag;ch[1]->tag+=tag;
        key+=tag;tag=0;
    }
}treap[maxn],nul;
typedef node* pnode;
pnode gen[maxn],len,null=&nul;
int tot,ans,n,K,lnk[maxn],nxt[maxe],son[maxe],w[maxe],c[maxn];
bool vis[maxn];
void rot(pnode &p,int d){
    pnode k=p->ch[d^1];p->ch[d^1]=k->ch[d];k->ch[d]=p;
    p->maintain();k->maintain();p=k;
}
pnode newnode(int key){
    len->ch[0]=len->ch[1]=null;len->key=key;len->cnt=len->s=1;len->f=rand();len->tag=0;
    return len++;
}
void insert(pnode &p,int key){
    p->pushdown();
    if(p==null)p=newnode(key);else
    if(p->key==key)p->cnt++,p->s++;else{
        int d=key>p->key;
        insert(p->ch[d],key);
        if(p->ch[d]->f<p->f)rot(p,d^1);
    }
    p->maintain();
}
int get_n(pnode p,int key){
    if(p==null)return 0;
    p->pushdown();
    if(p->key==key)return p->ch[0]->s+p->cnt;
    if(p->key>key)return get_n(p->ch[0],key);
             else return get_n(p->ch[1],key)+p->ch[0]->s+p->cnt;
}
void print(pnode p){
    p->pushdown();
    if(p==null)return;
    print(p->ch[0]);
    for(int i=1;i<=p->cnt;i++)c[++c[0]]=p->key;
    print(p->ch[1]);
}
pnode merge(pnode p1,pnode p2){
    c[0]=0;
    if(p1->s>p2->s)swap(p1,p2);
    print(p1);
    for(int i=1;i<=c[0];i++)ans+=get_n(p2,K-c[i]);
    for(int i=1;i<=c[0];i++)insert(p2,c[i]);
    return p2;
}
void print_Treap(pnode p){
    if(p==null)return;
    p->pushdown();
    print_Treap(p->ch[0]);
    for(int i=1;i<=p->cnt;i++)printf("%d ",p->key);
    print_Treap(p->ch[1]);
}
pnode dfs(int x){
    vis[x]=0;pnode p=newnode(0);
    for(int j=lnk[x];j;j=nxt[j]) if(vis[son[j]]){
        pnode p2=dfs(son[j]);
        p2->tag+=w[j];
        p=merge(p2,p);
    }
    return p;
}
void add(int x,int y,int z){
    nxt[++tot]=lnk[x];son[tot]=y;w[tot]=z;lnk[x]=tot;
}
int main(){
    scanf("%d%d",&n,&K);
    null->ch[0]=null->ch[1]=null;null->key=null->s=null->cnt=null->tag=0;
    while(n||K){
        tot=0;memset(lnk,0,sizeof(lnk));memset(vis,1,sizeof(vis));
        for(int i=1,x,y,z;i<n;i++)scanf("%d%d%d",&x,&y,&z),add(x,y,z),add(y,x,z);
        ans=0;len=treap;
        dfs(1);
        printf("%d\n",ans);
        scanf("%d%d",&n,&K);
    }
    return 0;
}