红黑树及C++代码实现

红黑树及C++代码实现

红黑树是二叉搜索树的一种,单次插入、删除、查询的时间复杂度都是\(O(log(n))\)。红黑树的应用广泛,STL的set和map、Java的TreeSet和TreeMap等都是使用红黑树实现的

哨兵节点

在红黑树中,所有的叶子节点、根节点的父节点都是一个名为哨兵节点的节点。哨兵节点用于处理边界条件,可以很方便的识别树的边缘,用NIL来表示。哨兵节点替代的是原来空指针的作用,哨兵节点可以降低访问空指针的风险,同时也可以简化代码中的逻辑判断

性质

一颗红黑树必须满足下面五条性质:

  1. 节点为红色或黑色
  2. 根节点为黑色
  3. NIL节点为黑色
  4. 红色节点的子节点为黑色
  5. 从根节点到NIL节点的简单路径上的黑色节点数量相同(又或者说从任意一个节点出发到达NIL节点的所有简单路径所经过的黑色节点数量相同)

思考下面的二叉树是否为红黑树

        root(B)
          \
          N(R)
          / \
      NL(B) NR(B)

乍一看好像是红黑树,因为根节点到NL和NR的路径上的黑色节点数相同。但实际上这颗树并不是红黑树,将NIL节点画出来一目了然

        root(B)
        / \
      NIL  N(R)
          / \
      NL(B) NR(B)
      / \    / \
    NIL NIL NIL NIL

注意到根节点到达NIL节点有5条路径,经过的黑色节点(NIL算作黑色节点)数量分别为2、3、3、3、3,并不是全相等的,所以这颗二叉树不是红黑树

同时这也解释了为什么要使用哨兵节点

为什么红黑树效率高

从红黑树的性质出发,在不考虑NIL节点的情况下,假设根节点到叶节点的路径上经过的黑色节点数为\(r\),假设树的高度为\(h\),那么有

\[r \le h \le 2r \]

其实很好理解,在没有红色节点的情况下,路径上至少都会有\(r\)个黑色节点。有红色节点的情况下,路径节点最多的情况应该是黑红相间

假设红黑树中有\(n\)个节点(不包含NIL节点)那么有

\[n \ge 2^r - 1 \]

节点最少的情况下,节点应该全为黑色,此时红黑树是一颗满二叉树(红黑树的性质决定)

             root(B)
              / \
      N(R)           S(B)
      / \            / \
  NL(B) NR(B)     SL(B) SR(B)
  / \    / \      / \    / \
NIL NIL NIL NIL NIL NIL NIL NIL

上面的红黑树的高度为3(不计NIL节点),其节点数为\(2^3-1\),对于全黑的红黑树,其节点数为\(2^r - 1\)

而在全黑的情况下,我们还可以加上一些红色节点,也不会破坏红黑树的性质5,所以有\(n \ge 2^r - 1\)

做一个变形

\[r \le log_2(n + 1) \]

又因为\(r\le h \le 2r\)

所以可以得到树高\(h\),与节点数\(n\)的关系

\[log_2(n + 1) \le h \le 2log_2(n + 1) \]

红黑树又是一颗二叉搜索树,二叉搜索树进行插入、删除、查找的复杂度取决于树的最大高度,于是我们就可以知道红黑树进行插入、删除、查找的复杂度为\(O(log(n))\)

为什么使用红黑树而不是AVL树

AVL树的最大高度大约为\(\lfloor log_2(n) \rfloor + 1\),比红黑树少一个常数,那为什么STL和Java使用红黑树而不使用AVL树呢?

AVL树的高度小于红黑树,这让它在查询操作的表现比红黑树更加优秀,因为查询操作并不会改变树的结构。但是在面对大量数据写入和删除的情况,AVL树为了保持平衡性,会经常进行旋转,以保证AVL树的左子树和右子树的高度差不超过1。而红黑树对树结构的要求并不像AVL树那么严格,旋转的次数会少于AVL树。这就导致红黑树的整体性能比AVL树更加优秀

旋转

旋转操作用于调整树的结构,很常用

左旋

假如对P点左旋,那么将P的右节点B变为P的父节点,将P变为其左节点A的父节点,将P的右节点B的左节点BL变为P的右节点

       P                                                 B
      / \                     P点左旋                    / \
   A       B                  ----->                  P    BR
  / \     / \                                        / \
 AL AR   BL BR                                      A   BL
                                                   / \
                                                  AL AR

右旋

与左旋完全反过来

假如对P点右旋,那么将P的左节点A变为P的父节点,将P变为其右节点B的父节点,将P的左节点A的右节点AR变为P的左节点

       P                                                 A
      / \                     P点右旋                    / \
   A       B                  ----->                   AL  P
  / \     / \                                             / \
 AL AR   BL BR                                           AR  B
                                                            / \
                                                           BL BR

插入

插入节点颜色选取

根据红黑树的性质,节点必须为红色或者黑色中的一种。那么对于插入操作,插入的节点应该是什么颜色呢?

如果插入节点为黑色,那么我们将节点插入以后,红黑树的性质5一定会被打破,有一条路径上的黑色节点数会增加1,好像并不好调整,如果要调整,我们就需要将路径上的一个节点变为红色,以保证性质5成立,但是一个节点变为红色又可能会打破性质4

如果插入节点为红色,那么我们将节点插入以后,红黑树的性质5不会被打破,但性质4却也有可能会被打破

而在后面可以知道,连续两个红色节点是可以调整的

插入红色节点可以省去变红的操作,更加方便。因此选择红色作为插入节点的颜色

Case 1:树为空

直接插入节点,并将节点的颜色变为黑色

Case 2:父节点为黑色节点

直接插入,不需要额外操作

Case 3:父节点为红色节点

  • 叔叔节点为红色节点

    将父节点以及叔叔节点变为黑色,将祖先节点变为红色。但是调整后祖先节点的父节点可能为红色,违反了性质4,因此我们需要递归的去调整祖先节点,直到没有冲突出现

            G(B)                  G(R)
           /   \                 /   \
         P(R) U(R)   ----->   P(B)   U(B)
         /                   /
       N(R)               N(R)
    
  • 叔叔节点为黑色,且插入节点与父节点同为左节点或右节点

    这种情况可以直接调整,将祖先节点G旋转,方向与父节点P方向相反,旋转后将父节点变成黑色,祖先节点变为红色

         G(B)                 P(B)
         / \         G左旋    /   \
       U(B) P(R)    ----->  G(R)   N(R)
             \             / 
              N(R)       U(B)
    
  • 叔叔节点为黑色,且插入节点与父节点方向相反

    这种情况可以转换为上面一种情况,从而直接进行调整。先将父节点P旋转,旋转方向与插入节点的方向相反,旋转后即可按照上面一种情况调整

         G(B)                G(B)                 N(B)
        /  \       P右旋     /  \       G左旋     /   \
      U(B) P(R)   -----> U(B)   N(R)   ----->  G(R)   P(R)
           /                     \
         N(R)                    P(R)
    

删除

双黑节点

在进行删除操作之前,有必要先了解一下双黑节点是什么

在红黑树中,双黑节点(Double Black Node) 是一个逻辑标记,用于表示在删除操作后某个位置需要“补足”一个额外的黑色节点,以维护红黑树的黑高一致性(所有路径的黑色节点数量相同)。它本身不是一种真实的颜色,而是修复红黑树性质时的一种临时状态。

从上面的解释中可以知道,如果某节点被打上双黑节点标记,就意味着这个节点会额外算作一个黑色节点,以保证红黑树的性质。但一个节点不能扮演两个角色,我们需要通过旋转、变色等操作将双黑节点标记消除(注意双黑节点是一个逻辑标记),消除后整棵树仍然是红黑树。

在删除操作时,如果出现了双黑节点,不会立即调整树的结构,而是在之后的删除调整操作进行统一的操作

近侄节点与远侄节点

  • 近侄节点:兄弟节点靠近删除节点的子节点

  • 远侄节点:兄弟节点远离删除节点的子节点

      P
     / \
    N   S
       / \
      SL SR

现在有上面一棵树,对于节点N来说,S为N的兄弟节点,SL则为N的近侄节点,SR则为N的远侄节点

Case 1:删除节点为树中唯一节点

若删除节点为树中的唯一节点,直接删除即可

Case 2:删除节点没有非NIL的子节点

  • 删除节点为红色节点

    如果删除的节点为红色,直接删除即可,删除后的树依旧是一颗红黑树,满足红黑树的五条性质

    	     P(?)                                         P(?)
             / \               删除后                      / \
        N(R)     S(?)         ------>                   NIL S(?)
        / \      / \                                        / \
      NIL NIL SL(?) SR(?)                                SL(?) SR(?)
    
  • 删除的节点为黑色节点

    	     P(?)
             / \
        N(B)     S(?)
        / \      / \
      NIL NIL SL(?) SR(?)
    

    我们首先删除需要删除的节点N,补上一个NIL节点,所补的NIL节点会被打上双黑节点标记,我们用C来表示这个NIL节点

    如果C节点只算做一个黑色节点,那么肯定是不满足性质5的,但是如果C节点算作两个黑色节点,那么就满足性质5了。在后面的删除调整操作我们考虑如何将双黑节点标记去除掉

    	     P(?)                                     P(?)
             / \                                      / \
        N(B)     S(?)            ------>           C(DB) S(?)
        / \      / \                                     / \
      NIL NIL SL(?) SR(?)                            SL(?)  SR(?) 
    

Case 3:删除节点有且仅有一个非NIL的子节点

这种情况下删除节点只能是黑色,并且删除节点的非NIL子节点只能是红色

  1. 如果删除节点是红色,那么删除节点的非NIL子节点一定为黑色,只有这样才能满足性质4。但是注意,从删除节点出发,如果直接走NIL子节点方向,那么路径上经过的黑色节点数为1,但走非NIL子节点方向,路径上经过的黑色节点数是大于1的,至少有非NIL子节点和非NIL子节点下面的NIL子节点两个黑色节点,因此该情况不满足红黑树的性质。

        N(R)   
        / \    
      NIL NR(B)
      	  / \
        NIL NIL
    
  2. 如果删除节点是黑色,删除节点的非NIL节点也为黑色,如同第一种情况,从删除节点出发,走NIL节点方向,路径上经过2个黑色节点,而走非NIL节点方向,路径上至少经过3个黑色节点,因此该情况也不满足红黑树的性质

        N(B)   
        / \    
      NIL NR(B)
      	  / \
        NIL NIL
    

因此删除节点只能是黑色,删除节点的非NIL子节点只能是红色

对于这种情况,可以直接调整。删除节点,用非NIL子节点代替该节点,并将节点颜色变成黑色

	     P(?)                           P(?)
         /                              /
    N(B)                              NR(B)
    / \                               / \
  NIL NR(R)         ----->          NIL NIL
       / \
     NIL NIL

Case 4:删除节点有两个非NIL的子节点

这种情况我们需要找到删除节点的后继节点,即右子树中最小的节点,然后交换两个节点,颜色不交换,然后继续删除需要删除的节点

此时会发现,删除的情况变成了Case 2或者Case 3

这是因为二叉搜索树的后继节点没有左节点,因此后继节点最多只有一个非NIL节点

删除调整

删除调整操作的目的是为了去除在删除过程中出现的双黑节点标记,下面会分成几种情况进行讨论

Case 1: 兄弟节点为黑色且兄弟节点的两个子节点为黑色

  1. 如果父节点为红色

    只需要将兄弟节点S变成红色,父节点P变成黑色就可以直接去掉双黑节点

        P(R)                        P(B)
         / \                        / \
     C(DB)  S(B)      ----->      C(B) S(R)
              / \                      / \
           SL(B) SR(B)              SL(B) SR(B)
    

    调整前后,根节点经过P点到达NIL节点的所有简单路径经过的黑色节点数量是相同的

  2. 如果父节点为黑色

    还是先将兄弟节点S变成红色

    但是注意这里并不能像上面那样,直接将S变成红色然后去掉双黑节点就完了,虽然调整后在这颗子树上满足了性质5,但是在整棵树上,经过P节点的路径的黑色节点数会减少一个,与其他路径的黑色节点数量不同,不满足性质5

    因此我们需要为父节点打上双黑节点标记,如果P节点算作两个黑色节点,那么就满足性质5了,然后继续调整父节点P,直到没有双黑节点出现

         P(B)                       P(DB)
         / \                        / \
     C(DB)  S(B)       ----->    C(B) S(R)
              / \                      / \
          SL(B)  SR(B)             SL(B) SR(B) 
    

Case 2:兄弟节点为黑色并且远侄节点为红色

这种情况可以直接进行调整

这里以删除节点为父节点的左节点为例,将父节点P进行旋转(删除节点在哪边就往哪边旋转),将兄弟节点S的颜色变为父节点P的颜色,父节点P和远侄节点SR的颜色变为黑色,去掉双黑标记

        P(?)                           S(?)
        / \                            / \
     C(DB) S(B)        左旋          P(B)  SR(B) 
          / \         ----->        / \   
      SL(?) SR(R)                 C(B) SL(?) 

可以看到调整后无论是在局部,还是在整体都是满足红黑树的性质的

Case 3:兄弟节点为黑色并且近侄节点为红色,远侄点为黑色

这种情况需要转换为 Case 2 状态,进而进行调整

这里以删除节点为父节点的左节点为例,将兄弟节点S进行旋转(删除节点在哪边,往反方向旋转),将S的颜色变为红色,将SL的颜色变为黑色,调整后按照 Case 2 继续调整

      P(?)                                 P(?)
      / \                                  / \
   C(DB) S(B)          右旋             C(DB) SL(B) 
         / \          ----->                   \
     SL(R)  SR(B)                              S(R)
       	                                         \
       	                                        SR(B)

Case 4:兄弟节点为红色

根据性质4,父节点一定为黑色,兄弟节点的两个子节点一定是黑色

策略是将兄弟节点变成黑色,然后继续调整。

将父节点P进行旋转(删除节点在哪边就往哪边旋转),将兄弟节点S变成黑色,将父节点P变为红色。之后就可以按照前面的情况进行处理了

        P(B)                     S(B)
        / \                      / \
     C(DB) S(R)  ----->       P(R) SR(B)
           / \                / \
       SL(B) SR(B)         C(DB) SL(B)

此时SL作为新的兄弟节点,继续调整C(DB)节点

测试

正确性测试

  • 插入

    我们使用洛谷P3879来验证我们的插入是否有问题。经过测试可以AC这道题,通过代码放在最后面(有点长)

    • 删除

      删除没有找到什么简单的题目,于是就造了一个简单的数据,数据规模为\(10^7\)

      下面为生成数据的python代码

      import random
      
      in_path = "./data.in"
      out_path = "./data.out"
      # 数据大小
      N = int(1e7)
      # 数据范围
      minNum = 0
      # 最大范围我们取节点数的十倍
      maxNum = int(N * 10)
      # 删除节点为一般节点
      erase_num = N // 2
      
      if __name__ == '__main__':
          # 生成N个区间内不重复的随机数
          x = random.sample(range(minNum, maxNum), N)
          # 获取删除的下标
          idxs = random.sample(range(0, len(x)), erase_num)
          # 获取到需要删除的数
          y = [x[idx] for idx in idxs]
          # 我们要算出哪些还存在,需要进行集合运算
          set_y = set(y)
          set_x = set(x)
          # 计算出最终的答案
          set_y = list(set_x - set_y)
          # 输入
          with open(in_path, mode='w', encoding='utf-8') as f:
              f.write(f"{N}\n")
              for each in x:
                  f.write(f"{each}\n")
              f.write(f"{erase_num}\n")
              for each in y:
                  f.write(f"{each}\n")
          # 输出
          with open(out_path, mode='w', encoding='utf-8') as f:
              f.write(f"{len(set_y)}\n")
              for each in set_y:
                  f.write(f"{each}\n")
      

      下面为验证的python代码

      import subprocess
      import time
      from subprocess import *
      
      cpp_path = "./RBTree.cpp"
      exe_path = "./RBTree.exe"
      test_cpp_path = "./test.cpp"
      test_exe_path = "./test.exe"
      out_path = "./out.out"
      data_path = "./data.out"
      in_path = "./data.in"
      
      with open(in_path,"r",encoding="utf-8") as f :
          stdin = f.read()
      
      compile_res = subprocess.run(["g++",cpp_path,"-o",exe_path],capture_output=True,text=True,timeout=5)
      # 编译成功
      if compile_res.returncode == 0:
          st = time.time()
          run_res = subprocess.run([exe_path],input=stdin,capture_output=True,text=True,timeout=100)
          ed = time.time()
          print(f"running time: {(ed - st) * 1000.0}ms")
          out = run_res.stdout
          with open(out_path, "w", encoding="utf-8") as f:
              f.write(out)
              
      with open(data_path,"r",encoding="utf-8") as f1,open(out_path,"r",encoding="utf-8") as f2:
          N1 = int(f1.readline().strip())
          N2 = int(f2.readline().strip())
          if(N1 != N2) :
              print("Wrong Answer!")
          else:
              # 由于数据是不重复的且输出顺序并不固定,直接用集合判断
              data1 = set([int(x.strip()) for x in f1.readlines()])
              data2 = set([int(x.strip()) for x in f2.readlines()])
              if data1 == data2:
                  print("Accept!")
              else :
                  print("Wrong Answer!")
      

      下面为测试结果,没问题

    • 黑高检验

      我们还是来检测一下整颗红黑树的黑高是否正确,检测代码如下

      void checkBlackHeight()
      {
          bool flag = true;
          checkBlackHeightHelp(root, flag);
          if (flag)
              cout << "right!" << endl;
          else
              cout << "error!" << endl;
      }
      
      int checkBlackHeightHelp(Node *N, bool &flag)
      {
          // NIL节点为黑色
          if (N == NIL)
              return 1;
          int left = checkBlackHeightHelp(N->left, flag);
          int right = checkBlackHeightHelp(N->right, flag);
          if (left != right)
              flag = false;
          return left + (N->color == Color::BLACK ? 1 : 0);
      }
      

      我们在插入和删除后分别进行一次黑高检验,没有问题

性能测试

我们使用STL的set进行对比,数据还是上面删除检验的数据。但首先我们先对程序优化一下,主要是优化输入和输出,数据量大了cin和cout的速度很慢,关闭流同步速度会快很多

在主函数中加入下面的代码即可

ios::sync_with_stdio(false);
cin.tie(0);
cout.tie(0);

set测试代码如下

#include <iostream>
using namespace std;
#include <set>
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    set<int> s;
    int n;
    cin >> n;
    for (int i = 0; i < n; i++)
    {
        int x;
        cin >> x;
        s.insert(x);
    }
    int m;
    cin >> m;
    for (int i = 0; i < m; i++)
    {
        int x;
        cin >> x;
        s.erase(x);
    }
    cout << s.size() << endl;
    for (auto x : s)
        cout << x << endl;
}
  • RBTree测试结果

  • set测试结果

我们实现的红黑树快了将近10秒,还不错

代码实现

颜色表示

这里用枚举类型表示红色和黑色

enum Color
{
    RED,
    BLACK
};

红黑树节点表示

使用结构体来表示红黑树的节点

template <typename Key, typename Value>
struct RBTreeNode
{
    // 按照key值进行插入删除等操作
    Key key;
    // value为存储的数据
    Value value;
    // 左节点指针
    RBTreeNode<Key, Value> *left;
    // 右节点指针
    RBTreeNode<Key, Value> *right;
    // 父节点指针
    RBTreeNode<Key, Value> *parent;
    // 节点颜色
    Color color;
    RBTreeNode(Color c, RBTreeNode<Key, Value> *l = nullptr, RBTreeNode<Key, Value> *r = nullptr, RBTreeNode<Key, Value> *p = nullptr)
    {
        color = c, left = l, right = r, parent = p;
    }
    RBTreeNode(Key k, Value v, Color c, RBTreeNode<Key, Value> *l = nullptr, RBTreeNode<Key, Value> *r = nullptr, RBTreeNode<Key, Value> *p = nullptr)
    {
        key = k, value = v, color = c, left = l, right = r, parent = p;
    }
};

创建节点

为了避免直接New内存(因为不好看),写四个辅助函数用于创建红黑树的节点

但首先我们先给节点类型取个别名,原来节点类型太长了

typedef RBTreeNode<Key, Value> Node;
  • 创建红色空节点

    // 创建一个空的红色节点
    Node *createEmptyRedNode(Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(Color::RED, l, r, p);
    }
    
  • 创建红色非空节点

    // 创建一个非空红色节点
    Node *createRedNode(Key key, Value value, Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(key, value, Color::RED, l, r, p);
    }
    
  • 创建黑色空节点

    // 创建一个空的黑色节点
    Node *createEmptyBlackNode(Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(Color::BLACK, l, r, p);
    }
    
  • 创建黑色非空节点

    // 创建一个非空黑色节点
    Node *createBlackNode(Key key, Value value, Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(key, value, Color::BLACK, l, r, p);
    }
    

也为了避免直接用delete释放内存,同样写一个辅助函数来删除节点

// 销毁节点
Node *destroyNode(Node *node)
{
    delete node;
}

左旋

// 左旋操作
void leftRotate(Node *p)
{
    // 获取旋转节点的右儿子
    Node *rightSon = p->right;
    // 获取祖先节点
    Node *grandParent = p->parent;

    // 更新旋转节点
    p->right = rightSon->left;
    p->parent = rightSon;

    // 更新右节点
    // 如果右儿子的左节点不为NIL,则将右儿子的左节点的父节点设置为p
    if (rightSon->left != NIL)
        rightSon->left->parent = p;

    rightSon->parent = grandParent;
    rightSon->left = p;

    // 更新祖先节点
    // 如果p为根节点
    if (grandParent == NIL)
        root = rightSon;
    // 如果p为左节点
    else if (grandParent->left == p)
        grandParent->left = rightSon;
    // 如果p为右节点
    else
        grandParent->right = rightSon;
}

右旋

// 右旋操作
void rightRotate(Node *p)
{
    // 获取旋转节点的左儿子
    Node *leftSon = p->left;
    // 获取祖先节点
    Node *grandParent = p->parent;

    // 更新旋转节点
    p->left = leftSon->right;
    p->parent = leftSon;

    // 更新左节点
    // 如果左儿子的右节点不为NIL,则将左儿子的右节点的父节点设置为p
    if (leftSon->right != NIL)
        leftSon->right->parent = p;

    leftSon->parent = grandParent;
    leftSon->right = p;

    // 更新祖先节点
    // 如果p为根节点
    if (grandParent == NIL)
        root = leftSon;
    // 如果p为左节点
    else if (grandParent->left == p)
        grandParent->left = leftSon;
    // 如果p为右节点
    else
        grandParent->right = leftSon;
}

插入调整

// 插入调整
void insertFixup(Node *N)
{
    // 一直向上调整,直到父节点不是红色,或者父节点为根节点
    // 注意这里的N节点一定为红色
    while (N->parent->color == Color::RED && N->parent != root)
    {
        // 如果父节点为左子节点
        if (N->parent == N->parent->parent->left)
        {
            // 获取叔叔节点
            Node *U = N->parent->parent->right;
            // 如果叔叔节点的颜色也为红色,则将父节点和叔叔节点都变为黑色,祖先节点变为红色,继续向上调整祖先节点
            if (U->color == Color::RED)
            {
                U->color = Color::BLACK;
                N->parent->color = Color::BLACK;
                N->parent->parent->color = Color::RED;
                // 继续调整祖先节点
                N = N->parent->parent;
            }
            else
            {
                // 调整节点与父节点不同向,需旋转为同向
                if (N == N->parent->right)
                {
                    // 旋转N的父节点
                    N = N->parent;
                    // 旋转父节点
                    leftRotate(N);
                }
                // 调整节点与父节点同向
                // 注意要先变色再旋转,旋转之后会改变父子关系,到时候很混乱
                // 父节点设置为黑色
                N->parent->color = Color::BLACK;
                // 祖先节点设置为红色
                N->parent->parent->color = Color::RED;
                // 旋转祖先节点
                rightRotate(N->parent->parent);
            }
        }
        // 如果父节点为右子节点
        else
        {
            // 获取叔叔节点
            Node *U = N->parent->parent->left;
            // 如果叔叔节点的颜色也为红色,则将父节点和叔叔节点都变为黑色,祖先节点变为红色,继续向上调整祖先节点
            if (U->color == Color::RED)
            {
                U->color = Color::BLACK;
                N->parent->color = Color::BLACK;
                N->parent->parent->color = Color::RED;
                // 继续调整祖先节点
                N = N->parent->parent;
            }
            else
            {
                // 调整节点与父节点不同向,需旋转为同向
                if (N == N->parent->left)
                {
                    // 旋转N的父节点
                    N = N->parent;
                    // 旋转父节点
                    rightRotate(N);
                }
                // 调整节点与父节点同向
                // 注意要先变色再旋转,旋转之后会改变父子关系,到时候很混乱
                // 父节点设置为红色
                N->parent->color = Color::BLACK;
                // 祖先节点设置为红色
                N->parent->parent->color = Color::RED;
                // 旋转祖先节点
                leftRotate(N->parent->parent);
            }
        }
    }
    // 根节点可以无责任变成黑色
    root->color = Color::BLACK;
}

插入节点

void insert(Key key, Value value)
{
    Node *N = createRedNode(key, value, NIL, NIL, NIL);
    // 查询节点的父节点
    Node *p = NIL;
    // 迭代查找临时节点
    Node *temp = root;

    // 树为空,直接插入
    if (root == NIL)
    {
        N->color = BLACK;
        root = N;
        size++;
        return;
    }

    while (temp != NIL)
    {
        p = temp;
        // 往左查找
        if (N->key < temp->key)
            temp = temp->left;
        // 往右查找
        else if (N->key > temp->key)
            temp = temp->right;
        // key值已经存在,则替换数据
        else
        {
            temp->value = value;
            return;
        }
    }

    size++;
    N->parent = p;
    // 如果key值比父节点小,则作为左子节点
    if (N->key < p->key)
        p->left = N;
    // 如果key值比父节点大,则作为右子节点
    else
        p->right = N;

    insertFixup(N);
}

节点替换

在删除操作中,我们会经常将一个节点替换为另外一个节点。为了方便,我们封装一个函数用于节点替换

注意在这个替换函数中,只是更新了父节点与替换节点的关系,子节点并没有更新

// 用v替换u
void replace(Node *u, Node *v)
{
    // 如果u为root
    if (u->parent == NIL)
        root = v;
    // u为左子节点
    else if (u->parent->left == u)
        u->parent->left = v;
    // u为右子节点
    else
        u->parent->right = v;

    v->parent = u->parent;
}

查询后继节点

在删除操作中,有些情况需要找到删除节点的后继节点,为了方便,依旧是封装一个函数来查询后继节点

后继节点为右子树中最小的一个节点

// 获取最小节点
Node *getMinNode(Node *p)
{
    while (p->left != NIL)
        p = p->left;
    return p;
}

删除调整

// 删除调整
void removeFixup(Node *C)
{
    while (C != root && C->color == BLACK)
    {
        // 如果C为左子节点
        if (C == C->parent->left)
        {
            // 获取兄弟节点
            Node *S = C->parent->right;
            // 如果兄弟节点为红色,我们需要通过旋转将兄弟节点变为黑色
            if (S->color == Color::RED)
            {
                // 先变色再旋转
                S->color = Color::BLACK;
                C->parent->color = Color::RED;
                leftRotate(C->parent);
                // 获取新的兄弟节点
                S = C->parent->right;
            }

            // 兄弟节点的左儿子和右儿子都为黑色
            if (S->left->color == Color::BLACK && S->right->color == Color::BLACK)
            {
                // 父节点为红色
                if (C->parent->color == Color::RED)
                {
                    C->parent->color = Color::BLACK;
                    S->color = Color::RED;
                    // 已经调整完毕,不需要再进行调整就将C设置为根节点
                    C = root;
                }
                // 父节点为黑色
                else
                {
                    S->color = Color::RED;
                    // 父节点成为双黑节点,继续调整父节点
                    C = C->parent;
                }
            }
            else
            {
                // 近侄节点为红色,远侄子节点为黑色,需要调整为远侄节点为红色
                if (S->right->color == Color::BLACK)
                {
                    // 先变色再旋转
                    S->color = Color::RED;
                    S->left->color = Color::BLACK;
                    rightRotate(S);
                    // 由于旋转改变了父子关系,所以重新获取一下兄弟节点
                    S = C->parent->right;
                }
                // 调整后远侄节点为红色
                S->color = C->parent->color;
                C->parent->color = Color::BLACK;
                S->right->color = Color::BLACK;
                leftRotate(C->parent);
                // 已经调整完毕,不需要再进行调整就将C设置为根节点
                C = root;
            }
        }
        // 如果C为右子节点
        else
        {
            // 获取兄弟节点
            Node *S = C->parent->left;
            // 如果兄弟节点为红色,我们需要通过旋转将兄弟节点变为黑色
            if (S->color == Color::RED)
            {
                // 先变色再旋转
                S->color = Color::BLACK;
                C->parent->color = Color::RED;
                rightRotate(C->parent);
                // 获取新的兄弟节点
                S = C->parent->left;
            }

            // 兄弟节点的左儿子和右儿子都为黑色
            if (S->left->color == Color::BLACK && S->right->color == Color::BLACK)
            {
                // 父节点为红色
                if (C->parent->color == Color::RED)
                {
                    C->parent->color = Color::BLACK;
                    S->color = Color::RED;
                    // 已经调整完毕,不需要再进行调整就将C设置为根节点
                    C = root;
                }
                // 父节点为黑色
                else
                {
                    S->color = Color::RED;
                    // 父节点成为双黑节点,继续调整父节点
                    C = C->parent;
                }
            }
            else
            {
                // 近侄节点为红色,远侄子节点为黑色,需要调整为远侄节点为红色
                if (S->left->color == Color::BLACK)
                {
                    // 先变色再旋转
                    S->color = Color::RED;
                    S->right->color = Color::BLACK;
                    leftRotate(S);
                    // 由于旋转改变了父子关系,所以重新获取一下兄弟节点
                    S = C->parent->left;
                }
                // 调整后远侄节点为红色
                S->color = C->parent->color;
                C->parent->color = Color::BLACK;
                S->left->color = Color::BLACK;
                rightRotate(C->parent);
                // 已经调整完毕,不需要再进行调整就将C设置为根节点
                C = root;
            }
        }
    }
    // 根节点可以无责任变成黑色
    root->color = Color::BLACK;
}

删除节点

注意在删除节点有两个非NIL的子节点的情况下,需要交换节点,这里为了方便是直接交换的键和值的,但是如果键和值对象比较大,那么这样效率就很慢,最好还是通过指针交换来达到节点交换的目的

void remove(Key key)
{
    Node *N = root;
    while (N != NIL)
    {
        if (N->key == key)
            break;
        if (key < N->key)
            N = N->left;
        else
            N = N->right;
    }
    // 树中没有删除的key
    if (N == NIL)
        return;
    // 删除节点为树中唯一节点
    // Case 1
    if (size == 1)
    {
        if (root != NIL)
            destroyNode(root);
        root = NIL;
        size--;
        return;
    }
    // Case 4
    // 删除节点有两个非NIL的子节点
    if (N->left != NIL && N->right != NIL)
    {
        // 获得后继节点
        Node *minNode = getMinNode(N->right);
        // 我们这里直接交换键值,方便,但如果键值都是比较大的对象就很慢了,最好还是交换指针
        swap(minNode->key, N->key);
        swap(minNode->value, N->value);
        // 删除节点转换为后继节点,转移到case 2或 case 3
        N = minNode;
    }

    // Case 3
    if (N->left == NIL && N->right != NIL)
    {
        Node *rightSon = N->right;
        rightSon->color = N->color;
        // 用右子节点替换删除节点
        replace(N, rightSon);
        // 删除节点
        destroyNode(N);
    }
    // Case 3
    else if (N->left != NIL && N->right == NIL)
    {
        Node *leftSon = N->left;
        leftSon->color = N->color;
        // 用左子节点替换删除节点
        replace(N, leftSon);
        // 删除节点
        destroyNode(N);
    }
    // Case 2
    else
    {
        // 此情况为删除节点的两个儿子都为NIL节点
        // 删除节点为红色,直接删除即可
        if (N->color == Color::RED)
        {
            Node *parent = N->parent;
            if (parent->left == N)
                parent->left = NIL;
            else
                parent->right = NIL;
            destroyNode(N);
        }
        // 删除节点为黑色,出现双黑节点,需要向上调整
        else
        {
            removeFixup(N);
            // 调整后删除该节点
            Node *parent = N->parent;
            if (parent->left == N)
                parent->left = NIL;
            else
                parent->right = NIL;
            destroyNode(N);
        }
    }
    size--;
}

查询

查询操作也是很基本也很常用的操作,红黑树的查询操作与其他的二叉搜索树基本一样:比节点小就去左子树查,比节点大就去右子树查,和节点一样大说明查询到了

pair<bool, Value> query(Key key)
{
    if (root == NIL)
        return {false, Value()};
    Node *temp = root;

    while (temp != NIL)
    {
        if (key < temp->key)
            temp = temp->left;
        else if (key > temp->key)
            temp = temp->right;
        else
            return {true, temp->value};
    }
    // 最后没有查到
    return {false, Value()};
}

注意这个版本的查询返回的value与红黑树中的并不是同一个,而是一个拷贝。如果想通过返回的value操作树内的value,可以返回引用类型

pair<bool, Value&>

遍历数据

这里使用中序遍历整个红黑树,并统计红节点和黑节点个数

void print()
{
    int red = 0, black = 0;
    printOperation(root, red, black);
    cout << "redNode num = " << red << " " << "blackNode num = " << black << endl;
}
void printOperation(Node *p, int &red, int &black)
{
    if (p == NIL)
        return;
    printOperation(p->left, red, black);
    if (p->color == Color::RED)
        red++;
    else
        black++;
    cout << "key = " << p->key << " value = " << p->value << " color = " << p->color << endl;
    printOperation(p->right, red, black);
}

洛谷题目代码

下面代码与完整代码不同,只是实现了插入和查询操作

#include <iostream>
using namespace std;
#include <ctime>
#include <map>
#include <set>
#include <vector>
enum Color
{
    RED,
    BLACK
};
template <typename Key, typename Value>
struct RBTreeNode
{
    // 按照key值进行插入删除等操作
    Key key;
    // value为存储的数据
    Value value;
    // 左节点指针
    RBTreeNode<Key, Value> *left;
    // 右节点指针
    RBTreeNode<Key, Value> *right;
    // 父节点指针
    RBTreeNode<Key, Value> *parent;
    // 节点颜色
    Color color;
    RBTreeNode(Color c, RBTreeNode<Key, Value> *l = nullptr, RBTreeNode<Key, Value> *r = nullptr, RBTreeNode<Key, Value> *p = nullptr)
    {
        color = c, left = l, right = r, parent = p;
    }
    RBTreeNode(Key k, Value v, Color c, RBTreeNode<Key, Value> *l = nullptr, RBTreeNode<Key, Value> *r = nullptr, RBTreeNode<Key, Value> *p = nullptr)
    {
        key = k, value = v, color = c, left = l, right = r, parent = p;
    }
};

template <typename Key, typename Value>
class RBTree
{
    typedef RBTreeNode<Key, Value> Node;

private:
    // 根节点
    Node *root;
    // 哨兵节点
    Node *NIL;
    // 节点个数
    int size;
    // 创建一个空的红色节点
    Node *createEmptyRedNode(Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(Color::RED, l, r, p);
    }
    // 创建一个非空红色节点
    Node *createRedNode(Key key, Value value, Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(key, value, Color::RED, l, r, p);
    }
    // 创建一个空的黑色节点
    Node *createEmptyBlackNode(Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(Color::BLACK, l, r, p);
    }
    // 创建一个非空黑色节点
    Node *createBlackNode(Key key, Value value, Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(key, value, Color::BLACK, l, r, p);
    }
    // 销毁节点
    void destroyNode(Node *node)
    {
        delete node;
    }
    // 左旋操作
    void leftRotate(Node *p)
    {
        // 获取旋转节点的右儿子
        Node *rightSon = p->right;
        // 获取祖先节点
        Node *grandParent = p->parent;

        // 更新旋转节点
        p->right = rightSon->left;
        p->parent = rightSon;

        // 更新右节点
        // 如果右儿子的左节点不为NIL,则将右儿子的左节点的父节点设置为p
        if (rightSon->left != NIL)
            rightSon->left->parent = p;

        rightSon->parent = grandParent;
        rightSon->left = p;

        // 更新祖先节点
        // 如果p为根节点
        if (grandParent == NIL)
            root = rightSon;
        // 如果p为左节点
        else if (grandParent->left == p)
            grandParent->left = rightSon;
        // 如果p为右节点
        else
            grandParent->right = rightSon;
    }
    // 右旋操作
    void rightRotate(Node *p)
    {
        // 获取旋转节点的左儿子
        Node *leftSon = p->left;
        // 获取祖先节点
        Node *grandParent = p->parent;

        // 更新旋转节点
        p->left = leftSon->right;
        p->parent = leftSon;

        // 更新左节点
        // 如果左儿子的右节点不为NIL,则将左儿子的右节点的父节点设置为p
        if (leftSon->right != NIL)
            leftSon->right->parent = p;

        leftSon->parent = grandParent;
        leftSon->right = p;

        // 更新祖先节点
        // 如果p为根节点
        if (grandParent == NIL)
            root = leftSon;
        // 如果p为左节点
        else if (grandParent->left == p)
            grandParent->left = leftSon;
        // 如果p为右节点
        else
            grandParent->right = leftSon;
    }
    // 插入调整
    void insertFixup(Node *N)
    {
        // 一直向上调整,直到父节点不是红色,或者父节点为根节点
        // 注意这里的N节点一定为红色
        while (N->parent->color == Color::RED && N->parent != root)
        {
            // 如果父节点为左子节点
            if (N->parent == N->parent->parent->left)
            {
                // 获取叔叔节点
                Node *U = N->parent->parent->right;
                // 如果叔叔节点的颜色也为红色,则将父节点和叔叔节点都变为黑色,祖先节点变为红色,继续向上调整祖先节点
                if (U->color == Color::RED)
                {
                    U->color = Color::BLACK;
                    N->parent->color = Color::BLACK;
                    N->parent->parent->color = Color::RED;
                    // 继续调整祖先节点
                    N = N->parent->parent;
                }
                else
                {
                    // 调整节点与父节点不同向,需旋转为同向
                    if (N == N->parent->right)
                    {
                        // 旋转N的父节点
                        N = N->parent;
                        // 旋转父节点
                        leftRotate(N);
                    }
                    // 调整节点与父节点同向
                    // 注意要先变色再旋转,旋转之后会改变父子关系,到时候很混乱
                    // 父节点设置为黑色
                    N->parent->color = Color::BLACK;
                    // 祖先节点设置为红色
                    N->parent->parent->color = Color::RED;
                    // 旋转祖先节点
                    rightRotate(N->parent->parent);
                }
            }
            // 如果父节点为右子节点
            else
            {
                // 获取叔叔节点
                Node *U = N->parent->parent->left;
                // 如果叔叔节点的颜色也为红色,则将父节点和叔叔节点都变为黑色,祖先节点变为红色,继续向上调整祖先节点
                if (U->color == Color::RED)
                {
                    U->color = Color::BLACK;
                    N->parent->color = Color::BLACK;
                    N->parent->parent->color = Color::RED;
                    // 继续调整祖先节点
                    N = N->parent->parent;
                }
                else
                {
                    // 调整节点与父节点不同向,需旋转为同向
                    if (N == N->parent->left)
                    {
                        // 旋转N的父节点
                        N = N->parent;
                        // 旋转父节点
                        rightRotate(N);
                    }
                    // 调整节点与父节点同向
                    // 注意要先变色再旋转,旋转之后会改变父子关系,到时候很混乱
                    // 父节点设置为红色
                    N->parent->color = Color::BLACK;
                    // 祖先节点设置为红色
                    N->parent->parent->color = Color::RED;
                    // 旋转祖先节点
                    leftRotate(N->parent->parent);
                }
            }
        }
        // 根节点可以无责任变成黑色
        root->color = Color::BLACK;
    }
    // 用v替换u,只更换父节点关系
    void replace(Node *u, Node *v)
    {
        // 如果u为root
        if (u->parent == NIL)
            root = v;
        // u为左子节点
        else if (u->parent->left == u)
            u->parent->left = v;
        // u为右子节点
        else
            u->parent->right = v;

        v->parent = u->parent;
    }
    // 获取最小节点
    Node *getMinNode(Node *p)
    {
        while (p->left != NIL)
            p = p->left;
        return p;
    }

    // 查询全部元素辅助函数
    void queryAllHelp(Node *N, vector<Value> &v)
    {
        if (N == NIL)
            return;
        queryAllHelp(N->left, v);
        v.push_back(N->value);
        queryAllHelp(N->right, v);
    }

public:
    RBTree(/* args */)
    {
        NIL = createEmptyBlackNode();
        NIL->left = NIL, NIL->right = NIL, NIL->parent = NIL;
        root = NIL;
        size = 0;
    }
    pair<bool, Value &> query(Key key)
    {
        Value v;
        if (root == NIL)
            return {false, v};
        Node *temp = root;

        while (temp != NIL)
        {
            if (key < temp->key)
                temp = temp->left;
            else if (key > temp->key)
                temp = temp->right;
            else
                return {true, temp->value};
        }
        return {false, v};
    }
    void insert(Key key, Value value)
    {
        Node *N = createRedNode(key, value, NIL, NIL, NIL);
        // 查询节点的父节点
        Node *p = NIL;
        // 迭代查找临时节点
        Node *temp = root;

        // 树为空,直接插入
        if (root == NIL)
        {
            N->color = BLACK;
            root = N;
            size++;
            return;
        }

        while (temp != NIL)
        {
            p = temp;
            // 往左查找
            if (N->key < temp->key)
                temp = temp->left;
            // 往右查找
            else if (N->key > temp->key)
                temp = temp->right;
            // key值已经存在,则替换数据
            else
            {
                temp->value = value;
                return;
            }
        }

        size++;
        N->parent = p;
        // 如果key值比父节点小,则作为左子节点
        if (N->key < p->key)
            p->left = N;
        // 如果key值比父节点大,则作为右子节点
        else
            p->right = N;

        insertFixup(N);
    }

    int getSize()
    {
        return size;
    }
    vector<Value> queryAll()
    {
        vector<Value> temp;
        queryAllHelp(root, temp);
        return temp;
    }
    ~RBTree()
    {
    }
};

int main()
{
    RBTree<string, RBTree<int, int>> tr;
    int n;
    cin >> n;

    for (int i = 1; i <= n; i++)
    {
        int l;
        cin >> l;
        for (int j = 0; j < l; j++)
        {
            string s;
            cin >> s;
            pair<bool, RBTree<int, int> &> q = tr.query(s);
            if (q.first == false)
            {
                RBTree<int, int> temp;
                temp.insert(i, i);
                tr.insert(s, temp);
            }
            else
                q.second.insert(i, i);
        }
    }
    int m;
    cin >> m;
    for (int i = 0; i < m; i++)
    {
        string s;
        cin >> s;
        pair<bool, RBTree<int, int> &> q = tr.query(s);
        if (q.first == false)
            cout << endl;
        else
        {
            vector<int> ans = q.second.queryAll();
            for (auto x : ans)
                cout << x << " ";
            cout << endl;
        }
    }
}

完整代码

测试情形:第一行给定一个正整数n,随后输入n个数存储在红黑树,然后输入一个正整数m,随后输入m个数表示删除某个数。最后输出树中还存在多少数,并按照从小到大进行输出

#include <iostream>
using namespace std;
#include <ctime>
#include <map>
#include <set>
#include <vector>
enum Color
{
    RED,
    BLACK
};
template <typename Key, typename Value>
struct RBTreeNode
{
    // 按照key值进行插入删除等操作
    Key key;
    // value为存储的数据
    Value value;
    // 左节点指针
    RBTreeNode<Key, Value> *left;
    // 右节点指针
    RBTreeNode<Key, Value> *right;
    // 父节点指针
    RBTreeNode<Key, Value> *parent;
    // 节点颜色
    Color color;
    RBTreeNode(Color c, RBTreeNode<Key, Value> *l = nullptr, RBTreeNode<Key, Value> *r = nullptr, RBTreeNode<Key, Value> *p = nullptr)
    {
        color = c, left = l, right = r, parent = p;
    }
    RBTreeNode(Key k, Value v, Color c, RBTreeNode<Key, Value> *l = nullptr, RBTreeNode<Key, Value> *r = nullptr, RBTreeNode<Key, Value> *p = nullptr)
    {
        key = k, value = v, color = c, left = l, right = r, parent = p;
    }
};

template <typename Key, typename Value>
class RBTree
{
    typedef RBTreeNode<Key, Value> Node;

private:
    // 根节点
    Node *root;
    // 哨兵节点
    Node *NIL;
    // 节点个数
    int size;
    // 创建一个空的红色节点
    Node *createEmptyRedNode(Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(Color::RED, l, r, p);
    }
    // 创建一个非空红色节点
    Node *createRedNode(Key key, Value value, Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(key, value, Color::RED, l, r, p);
    }
    // 创建一个空的黑色节点
    Node *createEmptyBlackNode(Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(Color::BLACK, l, r, p);
    }
    // 创建一个非空黑色节点
    Node *createBlackNode(Key key, Value value, Node *l = nullptr, Node *r = nullptr, Node *p = nullptr)
    {
        return new Node(key, value, Color::BLACK, l, r, p);
    }
    // 销毁节点
    void destroyNode(Node *node)
    {
        delete node;
    }
    // 左旋操作
    void leftRotate(Node *p)
    {
        // 获取旋转节点的右儿子
        Node *rightSon = p->right;
        // 获取祖先节点
        Node *grandParent = p->parent;

        // 更新旋转节点
        p->right = rightSon->left;
        p->parent = rightSon;

        // 更新右节点
        // 如果右儿子的左节点不为NIL,则将右儿子的左节点的父节点设置为p
        if (rightSon->left != NIL)
            rightSon->left->parent = p;

        rightSon->parent = grandParent;
        rightSon->left = p;

        // 更新祖先节点
        // 如果p为根节点
        if (grandParent == NIL)
            root = rightSon;
        // 如果p为左节点
        else if (grandParent->left == p)
            grandParent->left = rightSon;
        // 如果p为右节点
        else
            grandParent->right = rightSon;
    }
    // 右旋操作
    void rightRotate(Node *p)
    {
        // 获取旋转节点的左儿子
        Node *leftSon = p->left;
        // 获取祖先节点
        Node *grandParent = p->parent;

        // 更新旋转节点
        p->left = leftSon->right;
        p->parent = leftSon;

        // 更新左节点
        // 如果左儿子的右节点不为NIL,则将左儿子的右节点的父节点设置为p
        if (leftSon->right != NIL)
            leftSon->right->parent = p;

        leftSon->parent = grandParent;
        leftSon->right = p;

        // 更新祖先节点
        // 如果p为根节点
        if (grandParent == NIL)
            root = leftSon;
        // 如果p为左节点
        else if (grandParent->left == p)
            grandParent->left = leftSon;
        // 如果p为右节点
        else
            grandParent->right = leftSon;
    }
    // 插入调整
    void insertFixup(Node *N)
    {
        // 一直向上调整,直到父节点不是红色,或者父节点为根节点
        // 注意这里的N节点一定为红色
        while (N->parent->color == Color::RED && N->parent != root)
        {
            // 如果父节点为左子节点
            if (N->parent == N->parent->parent->left)
            {
                // 获取叔叔节点
                Node *U = N->parent->parent->right;
                // 如果叔叔节点的颜色也为红色,则将父节点和叔叔节点都变为黑色,祖先节点变为红色,继续向上调整祖先节点
                if (U->color == Color::RED)
                {
                    U->color = Color::BLACK;
                    N->parent->color = Color::BLACK;
                    N->parent->parent->color = Color::RED;
                    // 继续调整祖先节点
                    N = N->parent->parent;
                }
                else
                {
                    // 调整节点与父节点不同向,需旋转为同向
                    if (N == N->parent->right)
                    {
                        // 旋转N的父节点
                        N = N->parent;
                        // 旋转父节点
                        leftRotate(N);
                    }
                    // 调整节点与父节点同向
                    // 注意要先变色再旋转,旋转之后会改变父子关系,到时候很混乱
                    // 父节点设置为黑色
                    N->parent->color = Color::BLACK;
                    // 祖先节点设置为红色
                    N->parent->parent->color = Color::RED;
                    // 旋转祖先节点
                    rightRotate(N->parent->parent);
                }
            }
            // 如果父节点为右子节点
            else
            {
                // 获取叔叔节点
                Node *U = N->parent->parent->left;
                // 如果叔叔节点的颜色也为红色,则将父节点和叔叔节点都变为黑色,祖先节点变为红色,继续向上调整祖先节点
                if (U->color == Color::RED)
                {
                    U->color = Color::BLACK;
                    N->parent->color = Color::BLACK;
                    N->parent->parent->color = Color::RED;
                    // 继续调整祖先节点
                    N = N->parent->parent;
                }
                else
                {
                    // 调整节点与父节点不同向,需旋转为同向
                    if (N == N->parent->left)
                    {
                        // 旋转N的父节点
                        N = N->parent;
                        // 旋转父节点
                        rightRotate(N);
                    }
                    // 调整节点与父节点同向
                    // 注意要先变色再旋转,旋转之后会改变父子关系,到时候很混乱
                    // 父节点设置为红色
                    N->parent->color = Color::BLACK;
                    // 祖先节点设置为红色
                    N->parent->parent->color = Color::RED;
                    // 旋转祖先节点
                    leftRotate(N->parent->parent);
                }
            }
        }
        // 根节点可以无责任变成黑色
        root->color = Color::BLACK;
    }
    // 删除调整
    void removeFixup(Node *C)
    {
        while (C != root && C->color == BLACK)
        {
            // 如果C为左子节点
            if (C == C->parent->left)
            {
                // 获取兄弟节点
                Node *S = C->parent->right;
                // 如果兄弟节点为红色,我们需要通过旋转将兄弟节点变为黑色
                if (S->color == Color::RED)
                {
                    // 先变色再旋转
                    S->color = Color::BLACK;
                    C->parent->color = Color::RED;
                    leftRotate(C->parent);
                    // 获取新的兄弟节点
                    S = C->parent->right;
                }

                // 兄弟节点的左儿子和右儿子都为黑色
                if (S->left->color == Color::BLACK && S->right->color == Color::BLACK)
                {
                    // 父节点为红色
                    if (C->parent->color == Color::RED)
                    {
                        C->parent->color = Color::BLACK;
                        S->color = Color::RED;
                        // 已经调整完毕,不需要再进行调整就将C设置为根节点
                        C = root;
                    }
                    // 父节点为黑色
                    else
                    {
                        S->color = Color::RED;
                        // 父节点成为双黑节点,继续调整父节点
                        C = C->parent;
                    }
                }
                else
                {
                    // 近侄节点为红色,远侄子节点为黑色,需要调整为远侄节点为红色
                    if (S->right->color == Color::BLACK)
                    {
                        // 先变色再旋转
                        S->color = Color::RED;
                        S->left->color = Color::BLACK;
                        rightRotate(S);
                        // 由于旋转改变了父子关系,所以重新获取一下兄弟节点
                        S = C->parent->right;
                    }
                    // 调整后远侄节点为红色
                    S->color = C->parent->color;
                    C->parent->color = Color::BLACK;
                    S->right->color = Color::BLACK;
                    leftRotate(C->parent);
                    // 已经调整完毕,不需要再进行调整就将C设置为根节点
                    C = root;
                }
            }
            // 如果C为右子节点
            else
            {
                // 获取兄弟节点
                Node *S = C->parent->left;
                // 如果兄弟节点为红色,我们需要通过旋转将兄弟节点变为黑色
                if (S->color == Color::RED)
                {
                    // 先变色再旋转
                    S->color = Color::BLACK;
                    C->parent->color = Color::RED;
                    rightRotate(C->parent);
                    // 获取新的兄弟节点
                    S = C->parent->left;
                }

                // 兄弟节点的左儿子和右儿子都为黑色
                if (S->left->color == Color::BLACK && S->right->color == Color::BLACK)
                {
                    // 父节点为红色
                    if (C->parent->color == Color::RED)
                    {
                        C->parent->color = Color::BLACK;
                        S->color = Color::RED;
                        // 已经调整完毕,不需要再进行调整就将C设置为根节点
                        C = root;
                    }
                    // 父节点为黑色
                    else
                    {
                        S->color = Color::RED;
                        // 父节点成为双黑节点,继续调整父节点
                        C = C->parent;
                    }
                }
                else
                {
                    // 近侄节点为红色,远侄子节点为黑色,需要调整为远侄节点为红色
                    if (S->left->color == Color::BLACK)
                    {
                        // 先变色再旋转
                        S->color = Color::RED;
                        S->right->color = Color::BLACK;
                        leftRotate(S);
                        // 由于旋转改变了父子关系,所以重新获取一下兄弟节点
                        S = C->parent->left;
                    }
                    // 调整后远侄节点为红色
                    S->color = C->parent->color;
                    C->parent->color = Color::BLACK;
                    S->left->color = Color::BLACK;
                    rightRotate(C->parent);
                    // 已经调整完毕,不需要再进行调整就将C设置为根节点
                    C = root;
                }
            }
        }
        // 根节点可以无责任变成黑色
        root->color = Color::BLACK;
    }
    // 用v替换u,只更换父节点关系
    void replace(Node *u, Node *v)
    {
        // 如果u为root
        if (u->parent == NIL)
            root = v;
        // u为左子节点
        else if (u->parent->left == u)
            u->parent->left = v;
        // u为右子节点
        else
            u->parent->right = v;

        v->parent = u->parent;
    }
    // 获取最小节点
    Node *getMinNode(Node *p)
    {
        while (p->left != NIL)
            p = p->left;
        return p;
    }
    // 打印辅助函数
    void printHelp(Node *p, int &red, int &black)
    {
        if (p == NIL)
            return;
        cout << "key = " << p->key << " value = " << p->value << " color = " << p->color << endl;
        printHelp(p->left, red, black);
        if (p->color == Color::RED)
            red++;
        else
            black++;

        printHelp(p->right, red, black);
    }
    // 查询全部元素辅助函数
    void queryAllHelp(Node *N, vector<Value> &v)
    {
        if (N == NIL)
            return;
        queryAllHelp(N->left, v);
        v.push_back(N->value);
        queryAllHelp(N->right, v);
    }

    int checkBlackHeightHelp(Node *N, bool &flag)
    {
        if (N == NIL)
            return 1;
        int left = checkBlackHeightHelp(N->left, flag);
        int right = checkBlackHeightHelp(N->right, flag);
        if (left != right)
            flag = false;
        return left + (N->color == Color::BLACK ? 1 : 0);
    }

public:
    RBTree(/* args */)
    {
        NIL = createEmptyBlackNode();
        NIL->left = NIL, NIL->right = NIL, NIL->parent = NIL;
        root = NIL;
        size = 0;
    }
    pair<bool, Value> query(Key key)
    {
        if (root == NIL)
            return {false, Value()};
        Node *temp = root;

        while (temp != NIL)
        {
            if (key < temp->key)
                temp = temp->left;
            else if (key > temp->key)
                temp = temp->right;
            else
                return {true, temp->value};
        }
        return {false, Value()};
    }
    void insert(Key key, Value value)
    {
        Node *N = createRedNode(key, value, NIL, NIL, NIL);
        // 查询节点的父节点
        Node *p = NIL;
        // 迭代查找临时节点
        Node *temp = root;

        // 树为空,直接插入
        if (root == NIL)
        {
            N->color = BLACK;
            root = N;
            size++;
            return;
        }

        while (temp != NIL)
        {
            p = temp;
            // 往左查找
            if (N->key < temp->key)
                temp = temp->left;
            // 往右查找
            else if (N->key > temp->key)
                temp = temp->right;
            // key值已经存在,则替换数据
            else
            {
                temp->value = value;
                return;
            }
        }

        size++;
        N->parent = p;
        // 如果key值比父节点小,则作为左子节点
        if (N->key < p->key)
            p->left = N;
        // 如果key值比父节点大,则作为右子节点
        else
            p->right = N;

        insertFixup(N);
    }

    void remove(Key key)
    {
        Node *N = root;
        while (N != NIL)
        {
            if (N->key == key)
                break;
            if (key < N->key)
                N = N->left;
            else
                N = N->right;
        }
        // 树中没有删除的key
        if (N == NIL)
            return;
        // 删除节点为树中唯一节点
        // Case 1
        if (size == 1)
        {
            if (root != NIL)
                destroyNode(root);
            root = NIL;
            size--;
            return;
        }
        // Case 4
        // 删除节点有两个非NIL的子节点
        if (N->left != NIL && N->right != NIL)
        {
            // 获得后继节点
            Node *minNode = getMinNode(N->right);
            // 我们这里直接交换键值,方便,但如果键值都是比较大的对象就很慢了,最好还是交换指针
            swap(minNode->key, N->key);
            swap(minNode->value, N->value);
            // 删除节点转换为后继节点,转移到case 2或 case 3
            N = minNode;
        }

        // Case 3
        if (N->left == NIL && N->right != NIL)
        {
            Node *rightSon = N->right;
            rightSon->color = N->color;
            // 用右子节点替换删除节点
            replace(N, rightSon);
            // 删除节点
            destroyNode(N);
        }
        // Case 3
        else if (N->left != NIL && N->right == NIL)
        {
            Node *leftSon = N->left;
            leftSon->color = N->color;
            // 用左子节点替换删除节点
            replace(N, leftSon);
            // 删除节点
            destroyNode(N);
        }
        // Case 2
        else
        {
            // 此情况为删除节点的两个儿子都为NIL节点
            // 删除节点为红色,直接删除即可
            if (N->color == Color::RED)
            {
                Node *parent = N->parent;
                if (parent->left == N)
                    parent->left = NIL;
                else
                    parent->right = NIL;
                destroyNode(N);
            }
            // 删除节点为黑色,出现双黑节点,需要向上调整
            else
            {
                removeFixup(N);
                // 调整后删除该节点
                Node *parent = N->parent;
                if (parent->left == N)
                    parent->left = NIL;
                else
                    parent->right = NIL;
                destroyNode(N);
            }
        }
        size--;
    }
    int getSize()
    {
        return size;
    }
    void print()
    {
        int red = 0, black = 0;
        printHelp(root, red, black);
        cout << "redNode num = " << red << " " << "blackNode num = " << black << endl;
    }

    vector<Value> queryAll()
    {
        vector<Value> temp;
        queryAllHelp(root, temp);
        return temp;
    }
    void checkBlackHeight()
    {
        bool flag = true;
        checkBlackHeightHelp(root, flag);
        if (flag)
            cout << "right!" << endl;
        else
            cout << "error!" << endl;
    }
    ~RBTree()
    {
    }
};

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    RBTree<int, int> tr;
    int n, m, x;
    cin >> n;
    for (int i = 1; i <= n; i++)
    {
        cin >> x;
        tr.insert(x, x);
    }
    // cout << "check black height after insert: ";
    // tr.checkBlackHeight();
    cin >> m;
    for (int i = 1; i <= m; i++)
    {
        cin >> x;
        tr.remove(x);
    }
    // cout << "check black height after remove: ";
    // tr.checkBlackHeight();
    vector<int> ans = tr.queryAll();
    cout << ans.size() << endl;
    for (auto x : ans)
        cout << x << endl;
}

感想

红黑树很复杂,特别是删除操作,要考虑很多情况,这篇文章也写了有几天,在网上也找了很多资料,终归是搞出来了,也是见识到了红黑树的”厉害“

水平有限,测试的情况比较少,代码不保证完全正确,如果有问题欢迎指正

posted @ 2025-04-22 19:33  zetxie  阅读(116)  评论(0)    收藏  举报