1. 程式人生 > 其它 >旋轉式Treap入門

旋轉式Treap入門

經典旋轉式Treap

Treap 是一種常見的平衡樹。二叉搜尋樹(BST)的缺陷是容易退化,而研究表明,隨機構造的 BST 是趨向於平衡的,Treap 就是一種更改了結點排序方式的 BST。Treap 為每個結點引入了一個額外的隨機權值 priority,一個結點的 priority 小於其子樹中所有結點的 priority(堆性質),這也是 Treap 這個名字的由來(Tree + Heap)。本文介紹旋轉式 Treap,還有一種無旋 Treap 以後再做介紹。

維持平衡

旋轉式 Treap 自然通過旋轉來維持平衡。不過,Treap 的旋轉只有兩種:zig 和 zag。

     A                  B        
    / \    --zig->     / \    
   B   C              D   A
  / \      <-zag--       / \  
 D   E                  E   C

當樹的形態發生變化,堆性質無法滿足時,就需要進行旋轉。

插入

先給結點隨機分配一個 priority,然後根據 BST 的性質將其插入到一個葉子上,再根據堆性質將該結點向上旋轉。

刪除

找到需要刪除的結點,向 priority 更大的兒子方向旋轉直到成為葉子,最後刪除。

查詢

和普通的 BST 類似。

實現

namespace treap {
std::mt19937 rnd(time(nullptr));
struct node {
    int size, key, priority, cnt, ch[2];
};
node t[maxn];
int tot = 0, root = 0;
void push_up(int p) {
    t[p].size = t[t[p].ch[0]].size + t[t[p].ch[1]].size + t[p].cnt;
}
/********************************
 *     A                  B     *
 *    / \    --zig->     / \    *
 *   B   C              D   A   *
 *  / \      <-zag--       / \  *
 * D   E                  E   C *
 ********************************/
void rotate(int &p, int op) {
    // op == 0 ? zig : zag
    int s = t[p].ch[op ^ 1];
    t[p].ch[op ^ 1] = t[s].ch[op];    // reconnect node E
    t[s].ch[op] = p, p = s;           // change father
    push_up(t[p].ch[op]), push_up(p); // update
}
void insert(int v, int &p = root) {
    if (!p) {
        p = ++tot;
        t[p].size = t[p].cnt = 1;
        t[p].key = v, t[p].priority = rnd();
        t[p].ch[0] = t[p].ch[1] = 0;
        return;
    }
    t[p].size++;
    if (t[p].key == v)
        t[p].cnt++;
    else if (t[p].key < v) {
        insert(v, t[p].ch[1]);
        if (t[t[p].ch[1]].priority < t[p].priority)
            rotate(p, 0); // zag
    } else {
        insert(v, t[p].ch[0]);
        if (t[t[p].ch[0]].priority < t[p].priority)
            rotate(p, 1); // zig
    }
}
void erase(int v, int &p = root) {
    if (!p)
        return;
    if (t[p].key == v) {
        if (t[p].cnt > 1) {
            t[p].cnt--;
            push_up(p);
            return;
        }
        if (!t[p].ch[0] || !t[p].ch[1]) {
            p = t[p].ch[0] + t[p].ch[1];
        } else if (t[t[p].ch[0]].priority < t[t[p].ch[0]].priority) {
            rotate(p, 1);
            erase(v, p);
        } else {
            rotate(p, 0);
            erase(v, p);
        }
    } else if (t[p].key < v) {
        erase(v, t[p].ch[1]);
        push_up(p);
    } else {
        erase(v, t[p].ch[0]);
        push_up(p);
    }
}
int rank(int v, int p = root) {
    if (!p)
        return 0;
    if (t[p].key == v)
        return t[t[p].ch[0]].size + 1;
    else if (t[p].key < v) {
        return t[t[p].ch[0]].size + t[p].cnt + rank(v, t[p].ch[1]);
    } else {
        return rank(v, t[p].ch[0]);
    }
}
int kth(int v, int p = root) {
    if (!p)
        return 0;
    if (v <= t[t[p].ch[0]].size) {
        return kth(v, t[p].ch[0]);
    } else if (v > t[t[p].ch[0]].size + t[p].cnt) {
        return kth(v - t[t[p].ch[0]].size - t[p].cnt, t[p].ch[1]);
    } else {
        return t[p].key;
    }
}
int pre(int v, int p = root) {
    int ret = inf;
    while (p) {
        if (t[p].key < v)
            ret = t[p].key, p = t[p].ch[1];
        else
            p = t[p].ch[0];
    }
    return ret;
}
int suf(int v, int p = root) {
    int ret = -inf;
    while (p) {
        if (t[p].key > v)
            ret = t[p].key, p = t[p].ch[0];
        else
            p = t[p].ch[1];
    }
    return ret;
}
}; // namespace treap

性質

旋轉式 Treap 的期望高度是 \(\Theta(\log n)\),並且在 insert 操作中旋轉的期望次數小於 \(2\)。和其他常見平衡樹相比,旋轉式 Treap 除了常數較小外並沒有什麼優勢。