Treap树 模板

Treap是一种非常好写而且易于理解的平衡树,然而效率基本看脸,毕竟用到了随机算法。

期望情况下 插入 删除等操作的复杂度均为$$O(log n)$$

详情见注释

 

#include 

using std::cin;
using std::cout;
using std::endl;
using std::string;
using int_t = long long int;
//表示treap中的一个节点
struct Node {
    //该节点的值
    int_t value;
    //该节点的随机优先级,用于将treap组织成一个小根堆
    //但堆是完全平衡的 而二叉搜索树显然不容易完全平衡 所以priority仅用来使treap带有堆的性质
    //所以treap仍然是二叉搜索树
    //某个节点的priority必须必它的子节点的小 
    int_t priority;
    //同样值的个数,用于值会重复的情况
    int_t count = 1;
    //以该节点为根节点的子树的节点数(考虑重复值的情况)
    int_t size = 1;
    //左子节点
    Node* left = nullptr;
    //右子节点
    Node* right = nullptr;
    //父节点 注意 这个值只会在查询时更新
    Node* parent = nullptr;

    Node(int_t value) :
    value(value) {
        priority = rand();
    }
    //重新计算该节点的size值
    void maintain() {
        size = count;
        if (left) size += left->size;
        if (right) size += right->size;
    }
    //返回当前节点左子节点的size 即比当前节点小的节点的个数
    int_t leftSize() {
        if (left) return left->size;
        else return 0;
    }
};
Node* root = nullptr;
//在treap中搜索一个节点,返回节点的指针,如果没找到就返回nullptr
Node* lookup(Node* node, int_t value, Node* from = nullptr) {
    if (node == nullptr) {
        return nullptr;
    }
    //搜索的过程中顺便设置一下父节点
    node->parent = from;
    if (node->value == value) {
        return node;
    } else if (value < node->value) {
        return lookup(node->left, value, node);
    } else {
        return lookup(node->right, value, node);
    }
}
//把node的右子节点绕node左旋
void leftRotate(Node*& node) {
    //把node的右节点存起来
    Node* temp = node->right;
    //然后把node的右节点换成tmp的左节点
    node->right = temp->left;
    //把temp的左节点换成node
    temp->left = node;
    //然后重新算一下size值
    node->maintain();
    temp->maintain();
    node = temp;

}
//把node的左子节点绕node右旋

void rightRotate(Node* & node) {
    Node* temp = node->left;
    node->left = temp->right;
    temp->right = node;
    node->maintain();
    temp->maintain();
    node = temp;

}
//向以node为根节点的子树中插入value 并且维护最小堆
void insert(Node*& node, int_t value) {

    if (node == nullptr) {
        //如果node==nullptr就说明应该在这个位置插入节点
        node = new Node(value);
    } else {
        //向左子树插入
        if (value < node->value) {
            insert(node->left, value);
            //插入结束后,如果插入的节点破坏了最小堆的性质,就需要旋转一下,把优先级小的节点转到上面
            //这样做不会破坏二叉搜索树的性质
            if (node->left->priority > node->priority) {
                rightRotate(node);
            }

        } else if (value > node->value) {
            insert(node->right, value);
            if (node->right->priority > node->priority) {
                leftRotate(node);
            }

        } else {
            //特殊情况 如果value已经存在于树中 则直接在计数器+1
            node->count++;
        }
    }
    //重新算一下size
    if (node) node->maintain();
}
//从以node为根节点的子树中删除value

void remove(Node* & node, int_t value) {
//找到了要删除的节点
    if (node->value == value) {
        //特殊情况 如果有多个值相同的节点 那么计数器-1即可
        if (node->count > 1) {
            node->count--;
            
        } else if (node->left == nullptr) {
            //如果当前节点没有左子节点 那么就直接把右子节点接上来
            Node* bak = node->right;
            delete node;
            node = bak;
        } else if (node->right == nullptr) {
            Node* bak = node->left;
            delete node;
            node = bak;
        } else {
            //如果node有两棵子树,那么就把优先级小的那一棵旋转到node,然后在另一个子树(此时node在这个子树里)删除node
            //执行完多次旋转后,最后一定会到node只有一棵子树或者没有子树的情况
            if (node->left->priority < node->right->priority) {
                rightRotate(node);
                remove(node->right, value);
            } else {
                leftRotate(node);
                remove(node->left, value);
            }
        }
    } else {
        if (value < node->value) {
            remove(node->left, value);
        } else {
            remove(node->right, value);
        }
    }
    if (node) node->maintain();
}
//输出先序遍历
void preOrder(Node* node) {
    if (node == nullptr) return;
    cout << node->value << endl;
    preOrder(node->left);
    preOrder(node->right);
}
//输出中序遍历
void inOrder(Node* node) {
    if (node == nullptr) return;
    inOrder(node->left);
    cout << node->value << endl;

    inOrder(node->right);
}
//输出后序遍历
void postOrder(Node* node) {
    if (node == nullptr) return;
    postOrder(node->left);
    postOrder(node->right);
    cout << node->value << endl;
}
//求某个节点的前驱 即比这个节点小的最大的节点
//返回该节点指针 若不存在则返回nullptr
//需要保证value一定存在于树中
Node* pre(Node* node, int_t value) {
    node = lookup(node, value);
    if (node->left) {
        //如果node有左子节点 那么node的前驱一定是node的左子节点的最右边的节点
        node = node->left;
        while (node->right) node = node->right;
        return node;
    } else {
        //如果node没有左子节点 则顺着node一直往上找 直到找到node是某个节点的左子节点为止
        Node* parent = node->parent;
        while (parent && node != parent->right) {
            node = parent;
            parent = parent->parent;
        }
        return parent;
    }
}
//求节点的后继 与前驱完全对称 不再赘述
Node* post(Node* node, int_t value) {
    node = lookup(node, value);
    if (node->right) {
        node = node->right;
        while (node->left) node = node->left;
        return node;
    } else {
        Node* parent = node->parent;
        while (node && node != parent->left) {
            node = parent;
            parent = parent->parent;

        }
        return parent;
    }
}
//在以node为根节点的子树中查询value的排名
//要求value必须存在于树中
int_t rank(Node* node, int_t value) {
    if (node == nullptr) return 0;
    //如果value恰好等于node->value 就说明node的左子树的所有节点恰好比value小 node的右子树的节点恰好比value 大 value的名次就是node左子树的节点数
    if (value == node->value) {
        return node->leftSize();
    } else if (value < node->value) {
        //如果value比node->value小 那么就应该去node的左子树确定value的名次
        return rank(node->left, value);
    } else {
        //如果value比node大 那么value的名次就是 node左子树的大小+值为node->value的元素个数+value在右子树的排名
        return rank(node->right, value) + node->count + node->leftSize();
    }
}
//在以node为根的子树内找第k小的元素
Node* kth(Node* node, int_t k) {
    if (node == nullptr || k < 1 || k > node->size) return nullptr;
    //如果k恰好是node的名次
    //因为同一个元素可能有多个 所以需要这样写
    if (node->leftSize() + 1 <= k && k <= node->leftSize() + node->count) {
        return node;
    } else if (k <= node->leftSize()) {
        return kth(node->left, k);
    } else {
        return kth(node->right, k - node->leftSize() - node->count);
    }
}

int main() {
//    freopen("Z:/data.in", "r", stdin);
//    freopen("Z:/data.out", "w", stdout);

    int_t n;
    cin>>n;
    while (n--) {
        string opt;
        cin>>opt;
        if (opt == "insert" || opt == "1") {
            int_t x;
            cin>>x;
            insert(root, x);
        } else if (opt == "find") {
            int_t x;
            cin>>x;
            cout << lookup(root, x) << endl;
        } else if (opt == "remove" || opt == "2") {
            int_t x;
            cin>>x;
            remove(root, x);
        } else if (opt == "pre") {
            preOrder(root);
        } else if (opt == "in") {
            inOrder(root);
        } else if (opt == "post") {
            postOrder(root);
        } else if (opt == "prev" || opt == "5") {
            int_t x;
            cin>>x;
            insert(root, x);
            cout << pre(root, x)->value << endl;
            remove(root, x);
        } else if (opt == "next" || opt == "6") {
            int_t x;
            cin>>x;
            insert(root, x);
            cout << post(root, x)->value << endl;
            remove(root, x);
        } else if (opt == "rank" || opt == "3") {
            int_t x;
            cin>>x;
            insert(root, x);
            cout << rank(root, x) + 1 << endl;
            remove(root, x);
        } else if (opt == "kth" || opt == "4") {
            int_t x;
            cin>>x;
            Node* node = kth(root, x);
            if (node)
                cout << node->value << endl;
            else cout << 0 << endl;
        }
    }
}

评论

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

这个站点使用 Akismet 来减少垃圾评论。了解你的评论数据如何被处理