左偏红黑树实现

以下为洛谷P3369和LibreOJ#104普通平衡树的AC代码

洛谷P6136 【模板】普通平衡树（数据加强版）也已AC

#include <bits/stdc++.h>

template<class Key, class Compare=std::less<Key>>
class RB_Tree
{
private:
    enum class Node_Color
    {
        BLACK = 0, RED = 1
    };

    struct Node
    {
        Key key;
        Node *left{nullptr}, *right{nullptr}, *father{nullptr};
        Node_Color color;
        size_t node_cnt{1}, size;

        Node() = default;

        Node(Key key, Node_Color color, size_t size, Node *father)
                : key(key), color(color), size(size), father(father) {}
    } *root{nullptr};

    Compare cmp;
    using KeyType = Key;


    void Destroy_Tree(Node *root) const
    {
        if (root != nullptr)
        {
            Destroy_Tree(root->left);
            Destroy_Tree(root->right);
            delete root;
        }
    }

    bool is_red(const Node *u) const { return u == nullptr ? false : u->color == Node_Color::RED; }

    size_t size(const Node *u) const { return u == nullptr ? 0 : u->size; }

    Node *rotate_left(Node *u) const
    {
        // left rotate a red link
        //          <u>                   <v>
        //        /    \\               //    \
        //       *      <v>    ==>     <u>     *
        //             /   \          /   \
        //            *     *        *     *

        assert(is_red(u->right));
        Node *v = u->right;
        u->right = v->left;
        v->left = u;
        if (u->right)u->right->father = u;
        v->father = u->father;
        u->father = v;
        v->color = u->color;
        u->color = Node_Color::RED;
        v->size = u->size;
        u->size = size(u->left) + size(u->right) + u->node_cnt;
        return v;
    }

    Node *rotate_right(Node *u) const
    {
        // right rotate a red link
        //            <u>               <v>
        //          //    \           /    \\
        //         <v>     *   ==>   *      <u>
        //        /   \                    /   \
        //       *     *                  *     *

        assert(is_red(u->left));
        Node *v = u->left;
        u->left = v->right;
        v->right = u;
        if (u->left)u->left->father = u;
        v->father = u->father;
        u->father = v;
        v->color = u->color;
        u->color = Node_Color::RED;
        v->size = u->size;
        u->size = size(u->left) + size(u->right) + u->node_cnt;
        return v;
    }

    void Negate_Node_Color(Node *u) const
    {
        u->color = u->color == Node_Color::RED ? Node_Color::BLACK : Node_Color::RED;
    }

    void color_flip(Node *u) const
    {
        Negate_Node_Color(u);
        Negate_Node_Color(u->left);
        Negate_Node_Color(u->right);
    }

    Node *get_key_node(const KeyType &key) const;

    Node *insert(Node *root, const KeyType &key, Node *father) const;

    Node *fix_up(Node *root) const;

    const KeyType &get_min_key(Node *root) const;

    size_t get_min_node_cnt(Node *root) const;

    Node *delete_min_node(Node *root);

    Node *delete_arbitrary_node(Node *root, const KeyType &key);

    Node *move_red_left(Node *root);

    Node *move_red_right(Node *root);

    void Print_Tree(Node *root, int indent) const;

    Node *left_node(Node *u) const;//直接前驱

    Node *right_node(Node *u) const;//直接后继

public:

    ~RB_Tree() { Destroy_Tree(root); }

    size_t size() const { return root == nullptr ? 0 : root->size; };

    void insert(const KeyType &key);

    bool erase(const KeyType &key);

    size_t count(const KeyType &key) const;

    void Print_Tree() const//脑袋逆时针旋转90度看输出
    {
        Print_Tree(root, 0);
        std::cout << std::endl;
    }

    size_t get_rank(const KeyType &x) const;

    const KeyType &kth(size_t rank) const;

    const KeyType &predecessor(const KeyType &key);//前驱

    const KeyType &successor(const KeyType &key);//后继

};

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::get_key_node(const KeyType &key) const
{
    Node *u = root;
    while (u != nullptr)
    {
        if (key == u->key)return u;
        else if (cmp(key, u->key))u = u->left;
        else u = u->right;
    }
    return nullptr;
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::fix_up(Node *root) const
{
    if (is_red(root->right) && !is_red(root->left))
        root = rotate_left(root);
    if (is_red(root->left) && is_red(root->left->left))
        root = rotate_right(root);
    if (is_red(root->left) && is_red(root->right))
        color_flip(root);
    root->size = size(root->left) + size(root->right) + root->node_cnt;
    return root;
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::insert(Node *root, const KeyType &key, Node *father) const
{
    if (root == nullptr)return new Node(key, Node_Color::RED, 1, father);
    if (root->key == key)
        root->node_cnt++;
    else if (cmp(key, root->key))//key < root->key
        root->left = insert(root->left, key, root);
    else
        root->right = insert(root->right, key, root);
    return fix_up(root);
}

template<class Key, class Compare>
void RB_Tree<Key, Compare>::insert(const KeyType &key)
{
    root = insert(root, key, nullptr);
    root->color = Node_Color::BLACK;
}

template<class Key, class Compare>
size_t RB_Tree<Key, Compare>::count(const KeyType &key) const
{
    Node *u = root;
    while (u != nullptr)
    {
        if (key == u->key)return u->node_cnt;
        else if (cmp(key, u->key))u = u->left;
        else u = u->right;
    }
    return 0;
}

template<class Key, class Compare>
const typename RB_Tree<Key, Compare>::KeyType &RB_Tree<Key, Compare>::get_min_key(Node *root) const
{
    while (root->left != nullptr)
        root = root->left;
    return root->key;
}

template<class Key, class Compare>
size_t RB_Tree<Key, Compare>::get_min_node_cnt(Node *root) const
{
    while (root->left != nullptr)
        root = root->left;
    return root->node_cnt;
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::move_red_left(Node *root)
{
    color_flip(root);
    if (is_red(root->right->left))//根据原理此时貌似不会出现root->right为空指针的情况
    {
        root->right = rotate_right(root->right);
        root = rotate_left(root);
        color_flip(root);
    }
    return root;
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::move_red_right(Node *root)
{
    color_flip(root);
    if (is_red(root->left->left))
    {
        root = rotate_right(root);
        color_flip(root);
    }
    return root;
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::delete_min_node(Node *root)
{
    if (root->left == nullptr)
    {
        delete root;
        return nullptr;
    }
    if (!is_red(root->left) && !is_red(root->left->left))
        root = move_red_left(root);
    root->left = delete_min_node(root->left);
    return fix_up(root);
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::delete_arbitrary_node(Node *root, const KeyType &key)
{
    if (cmp(key, root->key))//key < root->key
    {
        if (!is_red(root->left) && !is_red(root->left->left))
            root = move_red_left(root);
        root->left = delete_arbitrary_node(root->left, key);
    }
    else
    {
        if (is_red(root->left))root = rotate_right(root);//向右走要保证root或者root->rc是红色节点
        if (key == root->key && root->right == nullptr)
        {
            delete root;
            return nullptr;
        }
        //否则以右子树里最小值替代
        if (!is_red(root->right) && !is_red(root->right->left))
            root = move_red_right(root);
        if (key == root->key)
        {
            root->key = get_min_key(root->right);
            root->node_cnt = get_min_node_cnt(root->right);
            root->right = delete_min_node(root->right);
        }
        else
            root->right = delete_arbitrary_node(root->right, key);
    }
    return fix_up(root);
}

template<class Key, class Compare>
bool RB_Tree<Key, Compare>::erase(const KeyType &key)
{
    Node *u = get_key_node(key);
    if (u == nullptr)return false;
    else if (u->node_cnt > 1)
    {
        u->node_cnt--;
        while (u)
        {
            u->size = size(u->left) + size(u->right) + u->node_cnt;
            u = u->father;
        }
        return true;
    }
    else
    {
        root = delete_arbitrary_node(root, key);
        if (root != nullptr)root->color = Node_Color::BLACK;
        return true;
    }
}

template<class Key, class Compare>
size_t RB_Tree<Key, Compare>::get_rank(const KeyType &key) const
{
    int rank = 1;
    Node *u = root;
    while (u != nullptr)
    {
        if (cmp(key, u->key))u = u->left;
        else
        {
            rank += size(u->left);
            if (key == u->key)return rank;
            rank += u->node_cnt;
            u = u->right;
        }
    }
    return rank;
}

template<class Key, class Compare>
const typename RB_Tree<Key, Compare>::KeyType &RB_Tree<Key, Compare>::kth(size_t rank) const
{
    Node *u = root;
    while (u != nullptr)
    {
        if (size(u->left) >= rank)
            u = u->left;
        else
        {
            rank -= size(u->left);
            if (rank <= u->node_cnt)return u->key;
            rank -= u->node_cnt;
            u = u->right;
        }
    }
    throw std::runtime_error("rank > size(root)");
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::left_node(Node *u) const
{
    if (u->left == nullptr)
    {
        while (u->father && u->father->left == u)
            u = u->father;
        u = u->father;
    }
    else
    {
        u = u->left;
        while (u->right)
            u = u->right;
    }
    return u;
}

template<class Key, class Compare>
typename RB_Tree<Key, Compare>::Node *RB_Tree<Key, Compare>::right_node(Node *u) const
{
    if (u->right == nullptr)
    {
        while (u->father && u->father->right == u)
            u = u->father;
        u = u->father;
    }
    else
    {
        u = u->right;
        while (u->left)
            u = u->left;
    }
    return u;
}

template<class Key, class Compare>
const typename RB_Tree<Key, Compare>::KeyType &RB_Tree<Key, Compare>::predecessor(const KeyType &key)
{
    Node *u = root, *final_node = nullptr;
    while (u != nullptr)
    {
        final_node = u;
        if (cmp(u->key, key))u = u->right;
        else u = u->left;//key == u->key向左走
    }
    if (cmp(final_node->key, key))
        return final_node->key;
    else
        return left_node(final_node)->key;//注意key == final_node->key情况
}

template<class Key, class Compare>
const typename RB_Tree<Key, Compare>::KeyType &RB_Tree<Key, Compare>::successor(const KeyType &key)
{
    Node *u = root, *final_node = nullptr;
    while (u != nullptr)
    {
        final_node = u;
        if (cmp(key, u->key))u = u->left;
        else u = u->right;//key == u->key向右走
    }
    if (cmp(key, final_node->key))
        return final_node->key;
    else
        return right_node(final_node)->key;
}

template<class Key, class Compare>
void RB_Tree<Key, Compare>::Print_Tree(Node *root, int indent) const
{
    if (root == nullptr)return;
    Print_Tree(root->right, indent + 4);
    std::cout << std::string(indent, '-') << root->key << '\n';
    Print_Tree(root->left, indent + 4);
}

int main()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(0);
    RB_Tree<int> st;
    int n;
    std::cin >> n;
    while (n--)
    {
        int op, x;
        std::cin >> op >> x;
        switch (op)
        {
            case 1:
                st.insert(x);
                break;
            case 2:
                st.erase(x);
                break;
            case 3:
                std::cout << st.get_rank(x) << '\n';
                break;
            case 4:
                std::cout << st.kth(x) << '\n';
                break;
            case 5:
                std::cout << st.predecessor(x) << '\n';
                break;
            case 6:
                std::cout << st.successor(x) << '\n';
                break;
            default:
                assert(0);
        }
    }
    return 0;
}