红黑树

enum Color {
  RED,
  BLACK
}
 
enum Direction {
  ROOT = 0,
  LEFT = -1,
  RIGHT = 1
}
 
class RBTreeNode<T = any> {
  public left: RBTreeNode | null = null
  public right: RBTreeNode | null = null
  public parent: RBTreeNode | null = null
  public color: Color = Color.RED
  public key: number
  public value: T
 
  constructor(key: RBTreeNode['key'], value: T) {
    this.key = key
    this.value = value
  }
 
  public direction() {
    if (this.parent == null) return Direction.ROOT
    return this.parent.left == this ? Direction.LEFT : Direction.RIGHT
  }
 
  public isRoot() {
    return this.direction() == Direction.ROOT
  }
 
  public isLeft() {
    return this.direction() == Direction.LEFT
  }
 
  public isRight() {
    return this.direction() == Direction.RIGHT
  }
 
  public isRed() {
    return this.color == Color.RED
  }
 
  public isBlack() {
    return this.color == Color.BLACK
  }
 
  public sibling() {
    if (this.isRoot()) return null
    return this.isLeft() ? this.parent!.right : this.parent!.left
  }
 
  public uncle() {
    if (this.isRoot()) return null
    return this.parent!.sibling()
  }
 
  public grandParent() {
    if (this.isRoot()) return null
    return this.parent!.parent
  }
}
 
class RBTree<T = never> {
  private count: number
  public root: RBTreeNode<T> | null
 
  constructor() {
    this.count = 0
    this.root = null
  }
 
  private maintainRelationship(node: RBTreeNode) {
    if (node.left != null) {
      node.left.parent = node
    }
    if (node.right != null) {
      node.right.parent = node
    }
  }
 
  //     |                       |
  //     2      l-rotate(2)      4
  //    / \     ==========>     / \
  //   1   4    <==========    2   5
  //      / \   r-rotate(4)   / \
  //     3   5               1   3
 
  // 左旋
  private rotateLeft(node: RBTreeNode | null) {
    if (node == null || node.right == null) return
 
    const direction = node.direction()
    const [parent, successor] = [node.parent, node.right]
 
    node.right = successor.left
    successor.left = node
 
    this.maintainRelationship(node)
    this.maintainRelationship(successor)
 
    switch (direction) {
      case Direction.ROOT:
        this.root = successor
        break
      case Direction.LEFT:
        parent!.left = successor
        break
      case Direction.RIGHT:
        parent!.right = successor
        break
    }
 
    successor.parent = parent
  }
 
  // 右旋
  private rotateRight(node: RBTreeNode | null) {
    if (node == null || node.left == null) return
 
    const direction = node.direction()
    const [parent, successor] = [node.parent, node.left]
 
    node.left = successor.right
    successor.right = node
 
    this.maintainRelationship(node)
    this.maintainRelationship(successor)
 
    switch (direction) {
      case Direction.ROOT:
        this.root = successor
        break
      case Direction.LEFT:
        parent!.left = successor
        break
      case Direction.RIGHT:
        parent!.right = successor
        break
    }
 
    successor.parent = parent
  }
 
  private rotateSameDirection(node: RBTreeNode | null, direction: Direction) {
    direction == Direction.LEFT ? this.rotateLeft(node) : this.rotateRight(node)
  }
 
  private rotateOppositeDirection(node: RBTreeNode | null, direction: Direction) {
    direction == Direction.LEFT ? this.rotateRight(node) : this.rotateLeft(node)
  }
 
  // 清空
  public clear() {
    this.count = 0
    this.root = null
  }
 
  // 添加后
  private maintainAfterAdd(node: RBTreeNode) {
    if (node.isRoot()) return
 
    if (node.parent!.isBlack()) return
 
    if (node.parent!.isRoot()) {
      node.parent!.color = Color.BLACK
      return
    }
 
    if (node.uncle()?.isRed()) {
      node.parent!.color = Color.BLACK
      node.uncle()!.color = Color.BLACK
      node.grandParent()!.color = Color.RED
      this.maintainAfterAdd(node.grandParent()!)
      return
    }
 
    if (node.direction() != node.parent!.direction()) {
      const parent = node.parent!
      node.isLeft() ? this.rotateRight(parent) : this.rotateLeft(parent)
      node = parent
    }
 
    node.parent!.isLeft() ? this.rotateRight(node.grandParent()) : this.rotateLeft(node.grandParent())
 
    node.parent!.color = Color.BLACK
    node.sibling()!.color = Color.RED
  }
 
  // 添加
  public add(key: RBTreeNode['key'], value: T, replace = true) {
    const dfs = (node: RBTreeNode<T> | null) => {
      if (node == null) {
        this.root = new RBTreeNode(key, value)
        return
      }
 
      if (key == node.key) {
        if (replace) node.value = value
        return
      }
 
      if (key < node.key) {
        if (node.left == null) {
          node.left = new RBTreeNode(key, value)
          node.left.parent = node
          this.maintainAfterAdd(node.left)
          this.count++
        } else {
          dfs(node.left)
        }
      } else {
        if (node.right == null) {
          node.right = new RBTreeNode(key, value)
          node.right.parent = node
          this.maintainAfterAdd(node.right)
          this.count++
        } else {
          dfs(node.right)
        }
      }
    }
    dfs(this.root)
  }
 
  // 删除后
  private maintainAfterRemove(node: RBTreeNode) {
    if (node.isRoot()) return
 
    const direction = node.direction()
    let sibling = node.sibling()
 
    if (sibling?.isRed()) {
      const parent = node.parent!
      this.rotateSameDirection(parent, direction)
      sibling.color = Color.BLACK
      parent.color = Color.RED
      sibling = node.sibling()
    }
 
    let closeNephew = direction == Direction.LEFT ? sibling?.left : sibling?.right
    let distantNephew = direction == Direction.LEFT ? sibling?.right : sibling?.left
    const closeNephewIsBlack = closeNephew == null || closeNephew.isBlack()
    const distantNephewIsBlack = distantNephew == null || distantNephew.isBlack()
 
    if (closeNephewIsBlack && distantNephewIsBlack) {
      sibling!.color = Color.RED
      if (node.parent != null && node.parent.isRed()) {
        node.parent.color = Color.BLACK
      } else {
        this.maintainAfterRemove(node.parent!)
      }
    } else {
      if (closeNephew != null && closeNephew.isRed()) {
        this.rotateOppositeDirection(sibling, direction)
        closeNephew.color = Color.BLACK
        sibling!.color = Color.RED
        sibling = node.sibling()
        closeNephew = direction == Direction.LEFT ? sibling?.left : sibling?.right
        distantNephew = direction == Direction.LEFT ? sibling?.right : sibling?.left
      }
      this.rotateSameDirection(node.parent, direction)
      sibling!.color = node.parent!.color
      node.parent!.color = Color.BLACK
      if (distantNephew != null) distantNephew.color = Color.BLACK
    }
  }
 
  // 删除
  public remove(key: RBTreeNode['key']) {
    const dfs = (node: RBTreeNode | null): boolean => {
      if (node == null) return false
 
      if (key != node.key) {
        if (key < node.key) {
          const left = node.left
          if (left != null && dfs(left)) {
            this.maintainRelationship(node)
            return true
          } else {
            return false
          }
        } else {
          const right = node.right
          if (right != null && dfs(right)) {
            this.maintainRelationship(node)
            return true
          } else {
            return false
          }
        }
      }
 
      if (this.count == 1) {
        this.clear()
        return true
      }
 
      if (node.left != null && node.right != null) {
        let [parent, successor] = [node, node.right]
        while (successor.left != null) {
          parent = successor
          successor = successor.left
        }
        node.key = successor.key
        node.value = successor.value
        ;[node, successor] = [successor, node]
        this.maintainRelationship(parent)
      }
 
      if (node.left == null && node.right == null) {
        if (node.isBlack()) this.maintainAfterRemove(node)
        if (node.isLeft()) node.parent!.left = null
        else node.parent!.right = null
      } else {
        let [parent, replacement] = [node.parent, node.left != null ? node.left : node.right]
        switch (node.direction()) {
          case Direction.ROOT:
            this.root = replacement
            break
          case Direction.LEFT:
            parent!.left = replacement
            break
          case Direction.RIGHT:
            parent!.right = replacement
            break
        }
        if (!node.isRoot()) replacement!.parent = parent
        if (node.isBlack()) {
          if (replacement!.isRed()) {
            replacement!.color = Color.BLACK
          } else {
            this.maintainAfterRemove(replacement!)
          }
        }
      }
 
      this.count--
      return true
    }
    return dfs(this.root)
  }
 
  // 数据量
  public size() {
    return this.count
  }
}