温馨提示×

实现一个红黑树迭代器:C++编程挑战

c++
小樊
82
2024-04-26 19:04:03
栏目: 编程语言

#include <iostream>
#include <stack>
#include <queue>

enum Color {
    RED,
    BLACK
};

struct Node {
    int data;
    Color color;
    Node* left;
    Node* right;
    Node* parent;

    Node(int data, Color color) : data(data), color(color), left(nullptr), right(nullptr), parent(nullptr) {}
};

class RedBlackTree {
private:
    Node* root;

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

    void insert(int data) {
        Node* new_node = new Node(data, RED);
        if (root == nullptr) {
            root = new_node;
            root->color = BLACK;
            return;
        }

        Node* current = root;
        Node* parent = nullptr;

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

        new_node->parent = parent;
        if (data < parent->data) {
            parent->left = new_node;
        } else {
            parent->right = new_node;
        }

        fix_insert(new_node);
    }

    void fix_insert(Node* node) {
        while (node != root && node->parent->color == RED) {
            if (node->parent == node->parent->parent->left) {
                Node* uncle = node->parent->parent->right;
                if (uncle != nullptr && uncle->color == RED) {
                    node->parent->color = BLACK;
                    uncle->color = BLACK;
                    node->parent->parent->color = RED;
                    node = node->parent->parent;
                } else {
                    if (node == node->parent->right) {
                        node = node->parent;
                        rotate_left(node);
                    }
                    node->parent->color = BLACK;
                    node->parent->parent->color = RED;
                    rotate_right(node->parent->parent);
                }
            } else {
                Node* uncle = node->parent->parent->left;
                if (uncle != nullptr && uncle->color == RED) {
                    node->parent->color = BLACK;
                    uncle->color = BLACK;
                    node->parent->parent->color = RED;
                    node = node->parent->parent;
                } else {
                    if (node == node->parent->left) {
                        node = node->parent;
                        rotate_right(node);
                    }
                    node->parent->color = BLACK;
                    node->parent->parent->color = RED;
                    rotate_left(node->parent->parent);
                }
            }
        }
        root->color = BLACK;
    }

    void rotate_left(Node* node) {
        Node* right_child = node->right;
        node->right = right_child->left;
        if (right_child->left != nullptr) {
            right_child->left->parent = node;
        }
        right_child->parent = node->parent;
        if (node->parent == nullptr) {
            root = right_child;
        } else if (node == node->parent->left) {
            node->parent->left = right_child;
        } else {
            node->parent->right = right_child;
        }
        right_child->left = node;
        node->parent = right_child;
    }

    void rotate_right(Node* node) {
        Node* left_child = node->left;
        node->left = left_child->right;
        if (left_child->right != nullptr) {
            left_child->right->parent = node;
        }
        left_child->parent = node->parent;
        if (node->parent == nullptr) {
            root = left_child;
        } else if (node == node->parent->right) {
            node->parent->right = left_child;
        } else {
            node->parent->left = left_child;
        }
        left_child->right = node;
        node->parent = left_child;
    }

    void inorder_traversal(Node* node, std::queue<int>& q) {
        if (node == nullptr) {
            return;
        }
        inorder_traversal(node->left, q);
        q.push(node->data);
        inorder_traversal(node->right, q);
    }

    std::queue<int> inorder() {
        std::queue<int> q;
        inorder_traversal(root, q);
        return q;
    }
};

class RedBlackTreeIterator {
private:
    std::queue<int> inorder_list;

public:
    RedBlackTreeIterator(RedBlackTree& tree) {
        inorder_list = tree.inorder();
    }

    int next() {
        int data = inorder_list.front();
        inorder_list.pop();
        return data;
    }

    bool hasNext() {
        return !inorder_list.empty();
    }
};

int main() {
    RedBlackTree tree;
    tree.insert(10);
    tree

0