Poj 1741——treap的启发式合并

Tree

Description
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.
Input
The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.
Output
For each test case output the answer on a single line.
题目大意
给定一棵树，求出满足距离不大于k的点对的个数。
解题思想
最开始的想法就是做n次DFS然后累计答案，但是显然超时。
于是乎，我们想到如果能维护一个数据结构能够快速的知道小于k-disy的x的个数（x,y是枚举的点对，disx表示x到x,y最近祖先的距离，要满足disx+disy<=k，就是disx<=k-disy）,马上就想到treap，treap中维护的数值显然是到最近公共祖先的距离，但是每次进行合并时又涉及到+w，比较麻烦，所以维护到根的距离比较方便（累计答案的时候注意加两倍的最近祖先到根的距离）。合并其实很简单，把小树往大树里放就可以了，启发式合并其实就是暴力的一个一个放。也许有人会问为什么不会超时，很明显，最坏的情况是两个一样大的树进行合并，所以对于每个节点最多就合并log(n)次，询问也是log(n)的，所以总复杂度是O(nlog^2(n))。操作时要注意先询问答案后合并，否则会算重。

#include
#include
#include
using namespace std;
const int maxn=10005,maxm=20005;
struct jz{
    int x,s,w,ran,l,r;
}a[maxn*14];
int lnk[maxn],nxt[maxm],son[maxm],w[maxm],ro[maxn],tot,n,K,m,ans;
bool vis[maxn];
void add(int x,int y,int z){nxt[++tot]=lnk[x];lnk[x]=tot;son[tot]=y;w[tot]=z;}
void Putdata(int k){a[k].s=a[a[k].l].s+a[a[k].r].s+a[k].w;}
void rturn(int &k){
    int t=a[k].l;a[k].l=a[t].r;a[t].r=k;
    a[t].s=a[k].s;Putdata(k);k=t;
}
void lturn(int &k){
    int t=a[k].r;a[k].r=a[t].l;a[t].l=k;
    a[t].s=a[k].s;Putdata(k);k=t;
}
void Insert(int &k,int x){
    if (k==0){k=++m;a[k].s=a[k].w=1;a[k].ran=rand();a[k].x=x;return;}
    a[k].s++;
    if (x==a[k].x) a[k].w++;else
    if (xif (a[a[k].l].ranelse{
        Insert(a[k].r,x);
        if (a[a[k].r].ranvoid Join(int &k1,int k2){
    if (k2==0) return;
    for (int i=1;i<=a[k2].w;i++) Insert(k1,a[k2].x);
    Join(k1,a[k2].l);Join(k1,a[k2].r);
}
int Asksum(int k,int x){
    if (k==0) return 0;
    if (x==a[k].x) return a[k].w+a[a[k].l].s;else
    if (xreturn Asksum(a[k].l,x);else
    return a[k].w+a[a[k].l].s+Asksum(a[k].r,x);
}
int Count(int k1,int k2,int x){
    if (k2==0) return 0;
    return a[k2].w*Asksum(k1,x-a[k2].x)+Count(k1,a[k2].l,x)+Count(k1,a[k2].r,x);
}
void DFS(int x,int dep){
    vis[x]=1;
    for (int j=lnk[x];j;j=nxt[j])if (!vis[son[j]]){
        DFS(son[j],dep+w[j]);
        if (a[ro[x]].s2*dep);Join(ro[x],ro[son[j]]);
    }
    ans+=Asksum(ro[x],K+dep);
    Insert(ro[x],dep);
}
int main(){
    freopen("exam.in","r",stdin);
    freopen("exam.out","w",stdout);
    while (1){
        scanf("%d%d",&n,&K);
        if (n==0&&K==0) return 0;
        memset(a,0,sizeof(a));
        memset(ro,0,sizeof(ro));
        memset(vis,0,sizeof(vis));
        memset(lnk,0,sizeof(lnk));
        tot=ans=m=0;int x,y,z;
        for (int i=1;iscanf("%d%d%d",&x,&y,&z);add(x,y,z);add(y,x,z);}
        DFS(1,0);
        printf("%d\n",ans);
    }
}