c++

实现C++中支持多线程访问的线程安全红黑树

小樊
87
2024-04-26 19:34:05
栏目: 编程语言

#include <iostream>
#include <mutex>
#include <thread>
#include <chrono>
#include <random>
#include <condition_variable>

enum Color {
    RED,
    BLACK
};

template <typename T>
class Node {
public:
    T key;
    Node<T>* left;
    Node<T>* right;
    Node<T>* parent;
    Color color;

    Node(T key) : key(key), left(nullptr), right(nullptr), parent(nullptr), color(RED) {}
};

template <typename T>
class RedBlackTree {
private:
    Node<T>* root;
    std::mutex m;

public:
    RedBlackTree() : root(nullptr) {}

    void insert(T key) {
        std::lock_guard<std::mutex> lock(m);
        Node<T>* newNode = new Node<T>(key);

        if (root == nullptr) {
            root = newNode;
            root->color = BLACK;
            return;
        }

        Node<T>* current = root;
        Node<T>* parent = nullptr;

        while (current != nullptr) {
            parent = current;
            if (key < current->key) {
                current = current->left;
            } else {
                current = current->right;
            }
        }

        newNode->parent = parent;
        if (key < parent->key) {
            parent->left = newNode;
        } else {
            parent->right = newNode;
        }

        fixViolation(newNode);
    }

    void printInOrder() {
        std::lock_guard<std::mutex> lock(m);
        printInOrderHelper(root);
        std::cout << std::endl;
    }

private:
    void fixViolation(Node<T>* newNode) {
        Node<T>* parent = nullptr;
        Node<T>* grandparent = nullptr;

        while (newNode != root && newNode->color != BLACK && newNode->parent->color == RED) {
            parent = newNode->parent;
            grandparent = parent->parent;

            if (parent == grandparent->left) {
                Node<T>* uncle = grandparent->right;

                if (uncle != nullptr && uncle->color == RED) {
                    grandparent->color = RED;
                    parent->color = BLACK;
                    uncle->color = BLACK;
                    newNode = grandparent;
                } else {
                    if (newNode == parent->right) {
                        rotateLeft(parent);
                        newNode = parent;
                        parent = newNode->parent;
                    }

                    rotateRight(grandparent);
                    std::swap(parent->color, grandparent->color);
                    newNode = parent;
                }
            } else {
                Node<T>* uncle = grandparent->left;

                if (uncle != nullptr && uncle->color == RED) {
                    grandparent->color = RED;
                    parent->color = BLACK;
                    uncle->color = BLACK;
                    newNode = grandparent;
                } else {
                    if (newNode == parent->left) {
                        rotateRight(parent);
                        newNode = parent;
                        parent = newNode->parent;
                    }

                    rotateLeft(grandparent);
                    std::swap(parent->color, grandparent->color);
                    newNode = parent;
                }
            }
        }

        root->color = BLACK;
    }

    void rotateLeft(Node<T>* newNode) {
        Node<T>* rightChild = newNode->right;
        newNode->right = rightChild->left;
        if (newNode->right != nullptr) {
            newNode->right->parent = newNode;
        }
        rightChild->parent = newNode->parent;
        if (newNode->parent == nullptr) {
            root = rightChild;
        } else if (newNode == newNode->parent->left) {
            newNode->parent->left = rightChild;
        } else {
            newNode->parent->right = rightChild;
        }
        rightChild->left = newNode;
        newNode->parent = rightChild;
    }

    void rotateRight(Node<T>* newNode) {
        Node<T>* leftChild = newNode->left;
        newNode->left = leftChild->right;
        if (newNode->left != nullptr) {
            newNode->left->parent = newNode;
        }
        leftChild->parent = newNode->parent;
        if (newNode->parent == nullptr) {
            root = leftChild;
        } else if (newNode == newNode->parent->left) {
            newNode->parent->left = leftChild;
        } else {
            newNode->parent->right = leftChild;
        }
        leftChild->right = newNode;
        newNode->parent = leftChild;
    }

    void printInOrderHelper(Node<T>* current) {
        if (current == nullptr) {
            return;
        }

        printInOrderHelper(current->left);
        std::cout << current->key << " ";
        printInOrderHelper(current->right);
    }
};

int main() {
    RedBlackTree<int> tree;

    std::vector<std::thread> threads;
    for (int i = 0;

0
看了该问题的人还看了