红黑树

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
  }
}

const drawTree = <T>(
  root: T | null,
  payload: {
    left: (node: T) => T | null
    right: (node: T) => T | null
    value: (node: T) => string
  }
) => {
  // 中序遍历并记录节点深度
  const nodes: Array<{ depth: number; node: T }> = []
  let maxDepth = 0
  const inorderTraversal = (node: T | null, depth: number) => {
    if (node == null) return null
    inorderTraversal(payload.left(node), depth + 1)
    nodes.push({ depth, node })
    maxDepth = Math.max(maxDepth, depth)
    inorderTraversal(payload.right(node), depth + 1)
  }
  inorderTraversal(root, 0)
 
  // 填入二维矩阵；每一层中间多加了一行空行，用以填写‘树枝’
  const matrix = Array.from<unknown, string[]>({ length: 2 * maxDepth + 1 }, () => Array(nodes.length).fill(''))
  const colsLen = Array.from<number>({ length: nodes.length })
  const nodesMap = new Map<number, Map<number, T>>()
  for (let i = 0; i < nodes.length; i++) {
    const { depth, node } = nodes[i]
    const content = payload.value(node)
    matrix[depth * 2][i] = content
    colsLen[i] = content.length
    nodesMap.set(i, new Map().set(depth * 2, node))
  }
 
  // 添加‘树枝’
  for (let i = 0; i < matrix.length; i++) {
    for (let j = 0; j < matrix[0].length; j++) {
      if (matrix[i][j] == '') {
        matrix[i][j] = ' '.repeat(colsLen[j])
        continue
      }
      const node = nodesMap.get(j)?.get(i)
      if (!node) continue
      if (payload.left(node) != null) {
        matrix[i + 1][j] = '┘'.padEnd(colsLen[j], ' ')
        for (let k = j - 1; k >= 0; k--) {
          if (matrix[i + 2][k] != '') {
            matrix[i + 1][k] = '┌'.padEnd(colsLen[k], '─')
            break
          }
          matrix[i + 1][k] = ''.padEnd(colsLen[k], '─')
        }
      }
      if (payload.right(node) != null) {
        matrix[i + 1][j] = matrix[i + 1][j] != '' ? '┴'.padEnd(colsLen[j], '─') : '└'.padEnd(colsLen[j], '─')
        for (let k = j + 1; k < matrix[0].length; k++) {
          if (matrix[i + 2][k] != '') {
            matrix[i + 1][k] = '┐'.padEnd(colsLen[k], ' ')
            break
          }
          matrix[i + 1][k] = ''.padEnd(colsLen[k], '─')
        }
      }
    }
  }
 
  console.log(matrix.map(v => v.join('')).join('\n'))
}
 
const rbtree = new RBTree<number>()
Array.from({ length: 50 }, (v, k) => rbtree.add(k, k))
drawTree<RBTreeNode>(rbtree.root, {
  left: node => node.left,
  right: node => node.right,
  value: node => (node.isBlack() ? 'B' : 'R') + String(node.value)
})
const removes = Array.from({ length: 20 }, () => (Math.random() * 20) >> 0)
console.log('\nremove:%s\n', removes.join(', '))
removes.forEach(k => rbtree.remove(k))
drawTree<RBTreeNode>(rbtree.root, {
  left: node => node.left,
  right: node => node.right,
  value: node => (node.isBlack() ? 'B' : 'R') + String(node.value)
})