红黑树

#include <iostream>

enum class Color {
  Red,
  Black
};

template <typename T>
struct Node {
  T key;
  Node<T>* parent;
  Node<T>* left;
  Node<T>* right;
  Color color;

  Node(T k) : key(k), parent(nullptr), left(nullptr), right(nullptr), color(Color::Red) {}
};

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

  // 左旋
  void leftRotate(Node<T>* x);
  // 右旋
  void rightRotate(Node<T>* x);
  // 插入修正
  void insertFixup(Node<T>* z);
  // 将子树 u 替换为子树 v
  void transplant(Node<T>* u, Node<T>* v);
  // 求 x 的子树的最小值点
  Node<T>* minimum(Node<T>* x);
  // 删除修正
  void deleteFixup(Node<T>* x);

 public:
  RedBlackTree() : root(nullptr) {}
  // 插入
  void insert(T key);
  // 删除
  void remove(T key);
  // 查找
  Node<T>* search(T key);
  // 中序遍历
  void printInOrder(Node<T>* node);
  void printInOrderTraversal();
};

template <typename T>
void RedBlackTree<T>::leftRotate(Node<T>* x) {
  Node<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 RedBlackTree<T>::rightRotate(Node<T>* x) {
  Node<T>* y = x->left;
  x->left = y->right;
  if (y->right != nullptr) {
    y->right->parent = x;
  }
  y->parent = x->parent;
  if (x->parent == nullptr) {
    root = y;
  } else if (x == x->parent->right) {
    x->parent->right = y;
  } else {
    x->parent->left = y;
  }
  y->right = x;
  x->parent = y;
}

template <typename T>
void RedBlackTree<T>::insertFixup(Node<T>* z) {
  while (z->parent != nullptr && z->parent->color == Color::Red) {
    if (z->parent == z->parent->parent->left) {
      Node<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(z);
        }
        z->parent->color = Color::Black;
        z->parent->parent->color = Color::Red;
        rightRotate(z->parent->parent);
      }
    } else {
      Node<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(z);
        }
        z->parent->color = Color::Black;
        z->parent->parent->color = Color::Red;
        leftRotate(z->parent->parent);
      }
    }
  }
  root->color = Color::Black;
}

template <typename T>
void RedBlackTree<T>::transplant(Node<T>* u, Node<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>
Node<T>* RedBlackTree<T>::minimum(Node<T>* x) {
  while (x->left != nullptr) {
    x = x->left;
  }
  return x;
}

template <typename T>
void RedBlackTree<T>::deleteFixup(Node<T>* x) {
  while (x != root && (x == nullptr || x->color == Color::Black)) {
    if (x == x->parent->left) {
      Node<T>* w = x->parent->right;
      if (w->color == Color::Red) {
        w->color = Color::Black;
        x->parent->color = Color::Red;
        leftRotate(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) {
          if (w->left != nullptr) {
            w->left->color = Color::Black;
          }
          w->color = Color::Red;
          rightRotate(w);
          w = x->parent->right;
        }
        w->color = x->parent->color;
        x->parent->color = Color::Black;
        if (w->right != nullptr) {
          w->right->color = Color::Black;
        }
        leftRotate(x->parent);
        x = root;  // 退出循环
      }
    } else {
      Node<T>* w = x->parent->left;
      if (w->color == Color::Red) {
        w->color = Color::Black;
        x->parent->color = Color::Red;
        rightRotate(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) {
          if (w->right != nullptr) {
            w->right->color = Color::Black;
          }
          w->color = Color::Red;
          leftRotate(w);
          w = x->parent->left;
        }
        w->color = x->parent->color;
        x->parent->color = Color::Black;
        if (w->left != nullptr) {
          w->left->color = Color::Black;
        }
        rightRotate(x->parent);
        x = root;
      }
    }
  }
  if (x != nullptr) {
    x->color = Color::Black;
  }
}

template <typename T>
void RedBlackTree<T>::insert(T key) {
  Node<T>* z = new Node<T>(key);
  Node<T>* y = nullptr;
  Node<T>* x = root;
  while (x != nullptr) {
    y = x;
    if (z->key < x->key) {
      x = x->left;
    } else {
      x = x->right;
    }
  }
  z->parent = y;
  if (y == nullptr) {
    root = z;
  } else if (z->key < y->key) {
    y->left = z;
  } else {
    y->right = z;
  }
  insertFixup(z);
}

template <typename T>
Node<T>* RedBlackTree<T>::search(T key) {
  Node<T>* current = root;
  while (current != nullptr && current->key != key) {
    if (key < current->key) {
      current = current->left;
    } else {
      current = current->right;
    }
  }
  return current;
}

template <typename T>
void RedBlackTree<T>::remove(T key) {
  Node<T>* z = search(key);
  if (z == nullptr) {
    return;
  }
  Node<T>* x;
  Node<T>* y = z;
  Color originalColor = y->color;
  if (z->left == nullptr) {
    x = z->right;
    transplant(z, z->right);
  } else if (z->right == nullptr) {
    x = z->left;
    transplant(z, z->left);
  } else {
    y = minimum(z->right);
    originalColor = y->color;
    x = y->right;
    if (y->parent == z) {
      if (x != nullptr) {
        x->parent = y;
      }
    } else {
      transplant(y, y->right);
      y->right = z->right;
      y->right->parent = y;
    }
    transplant(z, y);
    y->left = z->left;
    y->left->parent = y;
    y->color = z->color;
  }
  if (originalColor == Color::Black) {
    deleteFixup(x);
  }
  delete z;
}

template <typename T>
void RedBlackTree<T>::printInOrder(Node<T>* node) {
  if (node != nullptr) {
    printInOrder(node->left);
    std::cout << node->key << " ";
    printInOrder(node->right);
  }
}

template <typename T>
void RedBlackTree<T>::printInOrderTraversal() {
  printInOrder(root);
  std::cout << std::endl;
}

int main() {
  RedBlackTree<int> tree;
  tree.insert(10);
  tree.insert(20);
  tree.insert(30);
  tree.insert(15);
  tree.insert(5);
  std::cout << "In-order traversal: ";
  tree.printInOrderTraversal();

  tree.remove(20);

  std::cout << "In-order traversal after removing 20: ";
  tree.printInOrderTraversal();

  return 0;
}