MK
摩柯社区 - 一个极简的技术知识社区
AI 面试

C++ 实现红黑树深入讲解

2024-04-105.8k 阅读

红黑树简介

红黑树(Red - Black Tree)是一种自平衡二叉搜索树,它在计算机科学中有着广泛的应用,例如在 Linux 内核的进程调度、C++ STL 中的 mapset 容器等场景都有使用。红黑树的每个节点都带有一个颜色属性,要么是红色,要么是黑色。这种颜色标记以及一些特定的规则保证了红黑树在插入和删除操作后能自动调整平衡,使得树的高度始终保持在对数级别,从而保证了各种操作(插入、删除、查找)的时间复杂度为 $O(\log n)$,其中 $n$ 是树中节点的数量。

红黑树节点需要满足以下五条性质:

  1. 每个节点要么是红色,要么是黑色。
  2. 根节点是黑色。
  3. 每个叶节点(NIL 节点,空节点)是黑色的。
  4. 如果一个节点是红色的,则它的两个子节点都是黑色的(即不存在两个连续的红色节点)。
  5. 从任意一个节点到其每个叶节点的所有路径都包含相同数目的黑色节点(黑色高度一致)。

C++ 中红黑树节点的定义

在 C++ 中,我们首先需要定义红黑树的节点结构。节点结构除了存储数据之外,还需要记录颜色、父节点、左子节点和右子节点的信息。

// 定义红黑树节点颜色
enum class Color {
    RED,
    BLACK
};

// 红黑树节点模板类
template <typename T>
class RBNode {
public:
    T data;
    Color color;
    RBNode* left;
    RBNode* right;
    RBNode* parent;

    // 节点构造函数
    RBNode(const T& value) : data(value), color(Color::RED), left(nullptr), right(nullptr), parent(nullptr) {}
};

在上述代码中,我们使用 enum class 定义了颜色类型 Color,它有两个取值 REDBLACK。然后定义了模板类 RBNode,用于表示红黑树的节点。节点包含数据 data、颜色 color、左子节点指针 left、右子节点指针 right 和父节点指针 parent。构造函数用于初始化节点的数据并将颜色设置为红色,因为新插入的节点默认是红色,这样可以尽量少地破坏红黑树的性质。

红黑树的基本操作

左旋操作

左旋操作是红黑树调整平衡的重要操作之一。左旋操作会改变树的结构,它将一个节点与其右子节点的位置进行调整。

// 左旋操作
template <typename T>
void leftRotate(RBNode<T>*& root, RBNode<T>* x) {
    RBNode<T>* y = x->right;
    x->right = y->left;
    if (y->left != nullptr) {
        y->left->parent = x;
    }
    y->parent = x->parent;
    if (x->parent == nullptr) {
        root = y;
    } else if (x == x->parent->left) {
        x->parent->left = y;
    } else {
        x->parent->right = y;
    }
    y->left = x;
    x->parent = y;
}

在左旋操作中,我们首先保存 x 的右子节点 y。然后将 y 的左子节点连接到 x 的右子节点位置,并更新其 parent 指针。接着更新 yparent 指针,将 y 连接到 x 的父节点位置,并根据 x 是其父节点的左子节点还是右子节点来决定 y 连接到父节点的哪个位置。最后将 x 连接到 y 的左子节点位置,并更新 xparent 指针。

右旋操作

右旋操作与左旋操作类似,是左旋操作的镜像。它将一个节点与其左子节点的位置进行调整。

// 右旋操作
template <typename T>
void rightRotate(RBNode<T>*& root, RBNode<T>* y) {
    RBNode<T>* x = y->left;
    y->left = x->right;
    if (x->right != nullptr) {
        x->right->parent = y;
    }
    x->parent = y->parent;
    if (y->parent == nullptr) {
        root = x;
    } else if (y == y->parent->right) {
        y->parent->right = x;
    } else {
        y->parent->left = x;
    }
    x->right = y;
    y->parent = x;
}

右旋操作的步骤与左旋操作相反。首先保存 y 的左子节点 x,然后将 x 的右子节点连接到 y 的左子节点位置,并更新其 parent 指针。接着更新 xparent 指针,将 x 连接到 y 的父节点位置,并根据 y 是其父节点的右子节点还是左子节点来决定 x 连接到父节点的哪个位置。最后将 y 连接到 x 的右子节点位置,并更新 yparent 指针。

插入操作

红黑树的插入操作分为两个主要步骤:首先按照二叉搜索树的插入方式将新节点插入到合适的位置,然后通过调整颜色和旋转操作来恢复红黑树的性质。

// 插入操作
template <typename T>
void insert(RBNode<T>*& root, const T& value) {
    RBNode<T>* z = new RBNode<T>(value);
    RBNode<T>* y = nullptr;
    RBNode<T>* x = root;

    while (x != nullptr) {
        y = x;
        if (z->data < x->data) {
            x = x->left;
        } else {
            x = x->right;
        }
    }

    z->parent = y;
    if (y == nullptr) {
        root = z;
    } else if (z->data < y->data) {
        y->left = z;
    } else {
        y->right = z;
    }

    while (z != root && z->parent->color == Color::RED) {
        if (z->parent == z->parent->parent->left) {
            RBNode<T>* y = z->parent->parent->right;
            if (y != nullptr && y->color == Color::RED) {
                z->parent->color = Color::BLACK;
                y->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->right) {
                    z = z->parent;
                    leftRotate(root, z);
                }
                z->parent->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                rightRotate(root, z->parent->parent);
            }
        } else {
            RBNode<T>* y = z->parent->parent->left;
            if (y != nullptr && y->color == Color::RED) {
                z->parent->color = Color::BLACK;
                y->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->left) {
                    z = z->parent;
                    rightRotate(root, z);
                }
                z->parent->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                leftRotate(root, z->parent->parent);
            }
        }
    }
    root->color = Color::BLACK;
}

在插入操作中,首先按照二叉搜索树的方式找到插入位置,并将新节点 z 插入。新节点 z 初始颜色为红色。然后通过一个 while 循环来调整红黑树的性质。如果 z 的父节点是红色,且 z 不是根节点,就需要进行调整。这里分两种情况,即 z 的父节点是其祖父节点的左子节点还是右子节点。

在每种情况下,又分两种子情况。如果 z 的叔叔节点(父节点的兄弟节点)是红色,那么通过简单的颜色调整来恢复红黑树的性质。如果叔叔节点是黑色,则通过旋转操作(左旋或右旋)和颜色调整来恢复红黑树的性质。最后将根节点的颜色设置为黑色。

删除操作

红黑树的删除操作相对插入操作更为复杂。删除操作同样分为两个主要步骤:首先按照二叉搜索树的删除方式删除节点,然后通过调整颜色和旋转操作来恢复红黑树的性质。

// 找到节点 x 的后继节点
template <typename T>
RBNode<T>* successor(RBNode<T>* x) {
    if (x->right != nullptr) {
        RBNode<T>* y = x->right;
        while (y->left != nullptr) {
            y = y->left;
        }
        return y;
    } else {
        RBNode<T>* y = x->parent;
        while (y != nullptr && x == y->right) {
            x = y;
            y = y->parent;
        }
        return y;
    }
}

// 删除操作
template <typename T>
void deleteNode(RBNode<T>*& root, RBNode<T>* z) {
    RBNode<T>* x;
    RBNode<T>* y = z;
    Color y_original_color = y->color;
    if (z->left == nullptr) {
        x = z->right;
        transplant(root, z, z->right);
    } else if (z->right == nullptr) {
        x = z->left;
        transplant(root, z, z->left);
    } else {
        y = successor(z);
        y_original_color = y->color;
        x = y->right;
        if (y->parent == z) {
            if (x != nullptr) {
                x->parent = y;
            }
        } else {
            transplant(root, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        transplant(root, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    if (y_original_color == Color::BLACK) {
        while (x != root && (x == nullptr || x->color == Color::BLACK)) {
            if (x == x->parent->left) {
                RBNode<T>* w = x->parent->right;
                if (w->color == Color::RED) {
                    w->color = Color::BLACK;
                    x->parent->color = Color::RED;
                    leftRotate(root, x->parent);
                    w = x->parent->right;
                }
                if ((w->left == nullptr || w->left->color == Color::BLACK) &&
                    (w->right == nullptr || w->right->color == Color::BLACK)) {
                    w->color = Color::RED;
                    x = x->parent;
                } else {
                    if (w->right == nullptr || w->right->color == Color::BLACK) {
                        w->left->color = Color::BLACK;
                        w->color = Color::RED;
                        rightRotate(root, w);
                        w = x->parent->right;
                    }
                    w->color = x->parent->color;
                    x->parent->color = Color::BLACK;
                    w->right->color = Color::BLACK;
                    leftRotate(root, x->parent);
                    x = root;
                }
            } else {
                RBNode<T>* w = x->parent->left;
                if (w->color == Color::RED) {
                    w->color = Color::BLACK;
                    x->parent->color = Color::RED;
                    rightRotate(root, x->parent);
                    w = x->parent->left;
                }
                if ((w->right == nullptr || w->right->color == Color::BLACK) &&
                    (w->left == nullptr || w->left->color == Color::BLACK)) {
                    w->color = Color::RED;
                    x = x->parent;
                } else {
                    if (w->left == nullptr || w->left->color == Color::BLACK) {
                        w->right->color = Color::BLACK;
                        w->color = Color::RED;
                        leftRotate(root, w);
                        w = x->parent->left;
                    }
                    w->color = x->parent->color;
                    x->parent->color = Color::BLACK;
                    w->left->color = Color::BLACK;
                    rightRotate(root, x->parent);
                    x = root;
                }
            }
        }
        if (x != nullptr) {
            x->color = Color::BLACK;
        }
    }
    delete z;
}

// 辅助函数:用于替换节点 u 和 v
template <typename T>
void transplant(RBNode<T>*& root, RBNode<T>* u, RBNode<T>* v) {
    if (u->parent == nullptr) {
        root = v;
    } else if (u == u->parent->left) {
        u->parent->left = v;
    } else {
        u->parent->right = v;
    }
    if (v != nullptr) {
        v->parent = u->parent;
    }
}

在删除操作中,首先找到要删除的节点 z。如果 z 只有一个子节点或者没有子节点,直接用其子节点替换 z。如果 z 有两个子节点,则找到 z 的后继节点 y,用 y 替换 z,并将 z 的子节点连接到 y 上。

在删除节点后,如果被删除节点 y 的颜色是黑色,会破坏红黑树的性质,需要通过一个 while 循环来调整。这里同样分两种情况,即 xy 的子节点或者 y 被删除后留下的空位)是其父节点的左子节点还是右子节点。

在每种情况下,根据 x 的兄弟节点 w 的颜色以及 w 的子节点的颜色来进行相应的颜色调整和旋转操作,以恢复红黑树的性质。最后将 x 的颜色设置为黑色,并删除原来的节点 z

红黑树的遍历

红黑树作为一种二叉树,可以进行常见的遍历操作,如中序遍历、前序遍历和后序遍历。这里以中序遍历为例,展示如何遍历红黑树。

// 中序遍历
template <typename T>
void inorderTraversal(RBNode<T>* root) {
    if (root != nullptr) {
        inorderTraversal(root->left);
        std::cout << root->data << " ";
        inorderTraversal(root->right);
    }
}

中序遍历按照左子树、根节点、右子树的顺序访问节点。在上述代码中,通过递归方式实现了红黑树的中序遍历,依次输出节点的数据。

完整示例代码

下面是一个完整的红黑树示例代码,包含节点定义、基本操作、插入、删除和遍历等功能。

#include <iostream>

// 定义红黑树节点颜色
enum class Color {
    RED,
    BLACK
};

// 红黑树节点模板类
template <typename T>
class RBNode {
public:
    T data;
    Color color;
    RBNode* left;
    RBNode* right;
    RBNode* parent;

    // 节点构造函数
    RBNode(const T& value) : data(value), color(Color::RED), left(nullptr), right(nullptr), parent(nullptr) {}
};

// 左旋操作
template <typename T>
void leftRotate(RBNode<T>*& root, RBNode<T>* x) {
    RBNode<T>* y = x->right;
    x->right = y->left;
    if (y->left != nullptr) {
        y->left->parent = x;
    }
    y->parent = x->parent;
    if (x->parent == nullptr) {
        root = y;
    } else if (x == x->parent->left) {
        x->parent->left = y;
    } else {
        x->parent->right = y;
    }
    y->left = x;
    x->parent = y;
}

// 右旋操作
template <typename T>
void rightRotate(RBNode<T>*& root, RBNode<T>* y) {
    RBNode<T>* x = y->left;
    y->left = x->right;
    if (x->right != nullptr) {
        x->right->parent = y;
    }
    x->parent = y->parent;
    if (y->parent == nullptr) {
        root = x;
    } else if (y == y->parent->right) {
        y->parent->right = x;
    } else {
        y->parent->left = x;
    }
    x->right = y;
    y->parent = x;
}

// 插入操作
template <typename T>
void insert(RBNode<T>*& root, const T& value) {
    RBNode<T>* z = new RBNode<T>(value);
    RBNode<T>* y = nullptr;
    RBNode<T>* x = root;

    while (x != nullptr) {
        y = x;
        if (z->data < x->data) {
            x = x->left;
        } else {
            x = x->right;
        }
    }

    z->parent = y;
    if (y == nullptr) {
        root = z;
    } else if (z->data < y->data) {
        y->left = z;
    } else {
        y->right = z;
    }

    while (z != root && z->parent->color == Color::RED) {
        if (z->parent == z->parent->parent->left) {
            RBNode<T>* y = z->parent->parent->right;
            if (y != nullptr && y->color == Color::RED) {
                z->parent->color = Color::BLACK;
                y->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->right) {
                    z = z->parent;
                    leftRotate(root, z);
                }
                z->parent->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                rightRotate(root, z->parent->parent);
            }
        } else {
            RBNode<T>* y = z->parent->parent->left;
            if (y != nullptr && y->color == Color::RED) {
                z->parent->color = Color::BLACK;
                y->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->left) {
                    z = z->parent;
                    rightRotate(root, z);
                }
                z->parent->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                leftRotate(root, z->parent->parent);
            }
        }
    }
    root->color = Color::BLACK;
}

// 找到节点 x 的后继节点
template <typename T>
RBNode<T>* successor(RBNode<T>* x) {
    if (x->right != nullptr) {
        RBNode<T>* y = x->right;
        while (y->left != nullptr) {
            y = y->left;
        }
        return y;
    } else {
        RBNode<T>* y = x->parent;
        while (y != nullptr && x == y->right) {
            x = y;
            y = y->parent;
        }
        return y;
    }
}

// 删除操作
template <typename T>
void deleteNode(RBNode<T>*& root, RBNode<T>* z) {
    RBNode<T>* x;
    RBNode<T>* y = z;
    Color y_original_color = y->color;
    if (z->left == nullptr) {
        x = z->right;
        transplant(root, z, z->right);
    } else if (z->right == nullptr) {
        x = z->left;
        transplant(root, z, z->left);
    } else {
        y = successor(z);
        y_original_color = y->color;
        x = y->right;
        if (y->parent == z) {
            if (x != nullptr) {
                x->parent = y;
            }
        } else {
            transplant(root, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        transplant(root, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    if (y_original_color == Color::BLACK) {
        while (x != root && (x == nullptr || x->color == Color::BLACK)) {
            if (x == x->parent->left) {
                RBNode<T>* w = x->parent->right;
                if (w->color == Color::RED) {
                    w->color = Color::BLACK;
                    x->parent->color = Color::RED;
                    leftRotate(root, x->parent);
                    w = x->parent->right;
                }
                if ((w->left == nullptr || w->left->color == Color::BLACK) &&
                    (w->right == nullptr || w->right->color == Color::BLACK)) {
                    w->color = Color::RED;
                    x = x->parent;
                } else {
                    if (w->right == nullptr || w->right->color == Color::BLACK) {
                        w->left->color = Color::BLACK;
                        w->color = Color::RED;
                        rightRotate(root, w);
                        w = x->parent->right;
                    }
                    w->color = x->parent->color;
                    x->parent->color = Color::BLACK;
                    w->right->color = Color::BLACK;
                    leftRotate(root, x->parent);
                    x = root;
                }
            } else {
                RBNode<T>* w = x->parent->left;
                if (w->color == Color::RED) {
                    w->color = Color::BLACK;
                    x->parent->color = Color::RED;
                    rightRotate(root, x->parent);
                    w = x->parent->left;
                }
                if ((w->right == nullptr || w->right->color == Color::BLACK) &&
                    (w->left == nullptr || w->left->color == Color::BLACK)) {
                    w->color = Color::RED;
                    x = x->parent;
                } else {
                    if (w->left == nullptr || w->left->color == Color::BLACK) {
                        w->right->color = Color::BLACK;
                        w->color = Color::RED;
                        leftRotate(root, w);
                        w = x->parent->left;
                    }
                    w->color = x->parent->color;
                    x->parent->color = Color::BLACK;
                    w->left->color = Color::BLACK;
                    rightRotate(root, x->parent);
                    x = root;
                }
            }
        }
        if (x != nullptr) {
            x->color = Color::BLACK;
        }
    }
    delete z;
}

// 辅助函数:用于替换节点 u 和 v
template <typename T>
void transplant(RBNode<T>*& root, RBNode<T>* u, RBNode<T>* v) {
    if (u->parent == nullptr) {
        root = v;
    } else if (u == u->parent->left) {
        u->parent->left = v;
    } else {
        u->parent->right = v;
    }
    if (v != nullptr) {
        v->parent = u->parent;
    }
}

// 中序遍历
template <typename T>
void inorderTraversal(RBNode<T>* root) {
    if (root != nullptr) {
        inorderTraversal(root->left);
        std::cout << root->data << " ";
        inorderTraversal(root->right);
    }
}

int main() {
    RBNode<int>* root = nullptr;
    insert(root, 5);
    insert(root, 3);
    insert(root, 8);
    insert(root, 2);
    insert(root, 4);
    insert(root, 7);
    insert(root, 9);

    std::cout << "中序遍历红黑树: ";
    inorderTraversal(root);
    std::cout << std::endl;

    deleteNode(root, root->left->right);

    std::cout << "删除节点 4 后中序遍历红黑树: ";
    inorderTraversal(root);
    std::cout << std::endl;

    return 0;
}

main 函数中,我们创建了一个红黑树,并插入了一些节点,然后进行中序遍历。接着删除一个节点,再次进行中序遍历,以验证红黑树的插入和删除操作是否正确。

通过以上内容,我们深入讲解了 C++ 中红黑树的实现原理、基本操作以及具体代码示例,希望能帮助读者更好地理解和应用红黑树这一重要的数据结构。