树堆treap

BST最大的问题在于不能自平衡,最坏情况下会退化成链。在BST的基础上,引入一个权值属性rnd,要求二叉树所有节点的rnd满足堆的性质,通过该成员的随机性来保持平衡,这就是Treap。易见,Treap是一种随机性数据结构。

关于满足堆的性质,大根堆或小根堆都可以,这里选择大根堆。

结构定义

struct Node {
    int key, rnd;
    Node *son[2];
    Node(int x):key(x), rnd(rand()) {
        son[0] = son[1] = nullptr;
    }
    int cmp(int x) const {
        if (x == key) return -1;
        return x < key ? 0 : 1;
    }
};

旋转节点

树的平衡是通过旋转来实现的,以下是旋转操作,其中d取0表示左旋,取1表示右旋。

void Rotate(Node *&o, int d) {
    Node *k = o->son[d^1];
    o->son[d^1] = k->son[d];
    k->son[d] = o;
    o = k;
}

添加节点

void Insert(Node *&o, int x) {
    if (o) {
        int d = o->cmp(x);
        if (d == -1) return;
        Insert(o->son[d], x);
        if (o->son[d]->rnd > o->rnd)
            Rotate(o, d^1);
    } else o = new Node(x);
}

删除节点

void Delete(Node *&o, int x) {
    if (!o) return;
    int d = o->cmp(x);
    if (d == -1) {
        if (!o->son[0]) o = o->son[1];
        else if (!o->son[1]) o = o->son[0];
        else {
            int d2 = o->son[0]->rnd > o->son[1]->rnd ? 1 : 0;
            Rotate(o, d2);
            Delete(o->son[d2], x);
        }
    } else Delete(o->son[d], x);
}

查找节点

bool Exist(Node *o, int x) {
    if (!o) return false;
    int d = o->cmp(x);
    if (d == -1) return true;
    return Exist(o->son[d], x);
}

二分搜索

void LowerBound(Node *o, int x, int &ans) {
    if (!o) return;
    if (o->key >= x) ans = min(ans, o->key);
    int d = x < o->key ? 0 : 1;
    LowerBound(o->son[d], x, ans);
}
void UpperBound(Node *o, int x, int &ans) {
    if (!o) return;
    if (o->key > x) ans = min(ans, o->key);
    int d = x < o->key ? 0 : 1;
    UpperBound(o->son[d], x, ans);
}

第K小与排名

除了上述增删查等平衡树基本操作之外,Treap还可以实现查找第K小元素,以及给定key,问其在集合中的排名情况。

为了实现这两项功能,需要额外再增加个siz成员,用于标识以当前节点为根的子树的节点总数。

struct Node {
    int key, rnd, siz;
    Node *son[2];
    Node(int x):key(x),rnd(rand()),siz(1) {
        son[0] = son[1] = nullptr;
    }
    int cmp(int x) {
        if (x == key) return -1;
        return x < key ? 0 : 1;
    }
};

在旋转时需要对siz进行更新。

int Size(Node *o) {
    return o ? o->siz : 0;
}
void Maintain(Node *o) {
    o->siz = 1 + Size(o->son[0]) + Size(o->son[1]);
}
void Rotate(Node *&o, int d) {
    Node *k = o->son[d^1];
    o->son[d^1] = k->son[d];
    k->son[d] = o;
    Maintain(o); Maintain(k);
    o = k;
}

添加、删除以及查找的代码都不变。

第K小

int Select(Node *o, int k) {
    int t = 1 + Size(o->son[0]);
    if (k < t) return Select(o->son[0], k);
    if (k > t) return Select(o->son[1], k-t);
    return o->key;
}

排名

int Rank(Node *o, int x) {
    if (!o) return 0;
    if (x <= o->key) return Rank(o->son[0], x);
    return 1 + Size(o->son[0]) + Rank(o->son[1], x);
}

Rank()返回的是小于x的元素个数,1+Rank()即为实际排名。

Table of Contents