跳到主要内容

红黑树(RedBlackTree)

自平衡二叉搜索树之一，与 4 阶 B 树具有等价性

必须满足的 5 个性质​

  1. 节点只能是 Red 或 Black
  2. 根节点是 Black
  3. 叶子节点(遐想节点)是 Black
  4. Red 节点的子节点都是 Black
  • Red 节点的父节点都是 Black
  • 根节点到叶子节点的所有路径上不能有两个连续的 Red 节点
  1. 从任一节点到叶子节点的所有路径都包含相同数目的 Black 节点

核心代码​

import { BlanceBinarySearchTree } from './blanceBinarySearchTree';
import { BinaryTreeNode } from './binaryTree';
import { IBinarySearchTree } from './binarySearchTree';
export enum Color {
  RED,
  BLACK,
}
const getColor = (node: RBTreeNode<any> | null) => (node === null ? Color.BLACK : node.color);
const isRed = (node: RBTreeNode<any> | null) => getColor(node) === Color.RED;
const isBlack = (node: RBTreeNode<any> | null) => getColor(node) === Color.BLACK;

export interface IRBTree<T> extends IBinarySearchTree<T> {
  print: () => string;
  levelOrder: () => T[];
}
export class RBTreeNode<T> extends BinaryTreeNode<T> {
  private _color = Color.RED;
  get color() {
    return this._color;
  }
  get sibling() {
    return super.sibling as RBTreeNode<T>;
  }
  /* istanbul ignore next */
  constructor(
    public val: T,
    public parent: RBTreeNode<T> | null = null,
    public left: RBTreeNode<T> | null = null,
    public right: RBTreeNode<T> | null = null
  ) {
    super(val, parent, left, right);
  }
  setRed() {
    return this.setColor(Color.RED);
  }
  setBlack() {
    return this.setColor(Color.BLACK);
  }
  setColor(color: Color) {
    this._color = color;
    return this;
  }
  toString() {
    return super.toString() + (isRed(this) ? '(R)' : '');
  }
}
export class RBTree<T> extends BlanceBinarySearchTree<T> implements IRBTree<T> {
  protected createNode(
    val: T,
    parent: BinaryTreeNode<T> | null = null,
    left: BinaryTreeNode<T> | null = null,
    right: BinaryTreeNode<T> | null = null
  ): BinaryTreeNode<T> {
    return new RBTreeNode(
      val,
      parent as RBTreeNode<T> | null,
      left as RBTreeNode<T> | null,
      right as RBTreeNode<T> | null
    );
  }
  protected afterAdd(node: RBTreeNode<T>) {
    let parent = node.parent;
    // 如果父节点不存在，则添加的是根节点，直接染黑
    if (parent === null) {
      node.setBlack();
      return;
    }
    // 如果父节点是黑节点，则不做处理
    if (isBlack(parent)) return;
    const grandParent = parent.parent!.setRed();
    const uncle = parent.sibling;
    // 如果叔父节点是红色，则进行转换后上溢
    if (isRed(uncle)) {
      parent.setBlack();
      uncle.setBlack();
      this.afterAdd(grandParent);
      return;
    }
    // 判断旋转
    if (parent.isLeftChild) {
      if (node.isRightChild) {
        this.rotateLeft(parent);
        [node, parent] = [parent, node];
      }
      this.rotateRight(grandParent);
    } else {
      if (node.isLeftChild) {
        this.rotateRight(parent);
        [node, parent] = [parent, node];
      }
      this.rotateLeft(grandParent);
    }
    parent.setBlack();
    node.setRed();
  }
  protected afterRemove(node: RBTreeNode<T>) {
    // 如果删除的是红色节点，则染黑后(考虑节点可能是被删除节点的唯一子节点时)不做处理
    if (isRed(node)) {
      node.setBlack();
      return;
    }
    const parent = node.parent;
    // 如果删除的时根节点则不动
    if (parent === null) return;
    // 判断删除的节点是否时左子节点（考虑可能是下溢产生的删除和正常删除）
    const isLeftChild = parent.left === null || node.isLeftChild;
    let sibling = isLeftChild ? parent.right! : parent.left!;
    if (isLeftChild) {
      if (isRed(sibling)) {
        this.rotateLeft(parent);
        sibling.setBlack();
        parent.setRed();
        sibling = parent.right!;
      }
      if (isBlack(sibling.left) && isBlack(sibling.right)) {
        const blackState = isBlack(parent);
        sibling.setRed();
        parent.setBlack();
        blackState && this.afterRemove(parent);
      } else {
        if (isBlack(sibling.right)) {
          this.rotateRight(sibling);
          sibling = sibling.parent!;
        }
        this.rotateLeft(parent);
        sibling.setColor(parent.color);
        parent.setBlack();
        sibling.right?.setBlack();
      }
    } else {
      // 如果叔父节点是红色，则进行旋转使叔父节点染黑
      if (isRed(sibling)) {
        this.rotateRight(parent);
        sibling.setBlack();
        parent.setRed();
        sibling = parent.left!;
      }
      // 如果左右子树都为黑，则无真实子节点，则使父节点下溢
      if (isBlack(sibling.left) && isBlack(sibling.right)) {
        const blackState = isBlack(parent);
        sibling.setRed();
        parent.setBlack();
        blackState && this.afterRemove(parent);
      } else {
        // 如果兄弟节点的左节点是黑色，则进行旋转
        if (isBlack(sibling.left)) {
          this.rotateLeft(sibling);
          sibling = sibling.parent!;
        }
        this.rotateRight(parent);
        sibling.setColor(parent.color);
        sibling.left?.setBlack();
        parent.setBlack();
      }
    }
  }
}

核心代码​

/* istanbul ignore file */
import { repeat } from 'lodash';
import { IRBTree } from './rbTree';
const BLACK = false;
const RED = true;
type Color = boolean;
const isRed = (node: RBTreeNode2<any> | null) => node?.isRed ?? false;
const isBlack = (node: RBTreeNode2<any> | null) => node?.isBlack ?? true;
export class RBTreeNode2<T> {
  private _color: Color = RED;
  get color() {
    return this._color;
  }
  get isRed() {
    return this.color === RED;
  }
  get isBlack() {
    return this.color === BLACK;
  }
  get childSize() {
    return this.left === null && this.right === null
      ? 0
      : this.left === null && this.right !== null
      ? 1
      : this.left !== null && this.right === null
      ? 1
      : 2;
  }
  left: RBTreeNode2<T> | null = null;
  right: RBTreeNode2<T> | null = null;
  constructor(public val: T, public parent: RBTreeNode2<T> | null = null) {}
  toRed() {
    return this.toColor(RED);
  }
  toBlack() {
    return this.toColor(BLACK);
  }
  toColor(v: Color) {
    this._color = v;
    return this;
  }
  toString() {
    return `${this.val}【${this.isRed ? 'R' : 'B'}】`;
  }
}
export class RBTree2<T> implements IRBTree<T> {
  private _root: RBTreeNode2<T> | null = null;
  get root() {
    return this._root;
  }
  private _size = 0;
  get size() {
    return this._size;
  }
  get empty() {
    return this.size === 0;
  }
  constructor(private compare: (t1: T, t2: T) => number) {}
  clear() {
    this._root = null;
  }
  contains(v: T): boolean {
    return this.findNode(this.root, v)?.val !== null;
  }
  private findNode(root: RBTreeNode2<T> | null, v: T): RBTreeNode2<T> | null {
    if (root === null) return null;
    if (root.val === v) return root;
    if (this.compare(root.val, v) > 0) return this.findNode(root.left, v);
    else return this.findNode(root.right, v);
  }
  add(v: T): void {
    const oldNode = this.findNode(this.root, v);
    if (oldNode !== null) {
      oldNode.val = v;
      return;
    }
    if (this.root === null) {
      const node = new RBTreeNode2(v);
      this._root = node;
      this._size++;
      this.afterAdd(node);
      return;
    }
    let parent = this.root;
    let pos: 'left' | 'right' = 'left';
    while (parent) {
      const compare = this.compare(parent.val, v);
      if (compare > 0) {
        if (parent.left === null) break;
        parent = parent.left;
      } else {
        if (parent.right === null) {
          pos = 'right';
          break;
        }
        parent = parent.right;
      }
    }
    const node = new RBTreeNode2(v, parent);
    parent[pos] = node;
    this._size++;
    this.afterAdd(node);
  }
  afterAdd(node: RBTreeNode2<T>): void {
    let parent = node.parent;
    if (parent === null) {
      node.toBlack();
      return;
    }
    if (isBlack(parent)) return;
    const grand = parent.parent!;
    const sibling = grand.left === parent ? grand.right! : grand.left!;
    if (isRed(sibling)) {
      node.toRed();
      parent.toBlack();
      sibling.toBlack();
      this.afterAdd(grand.toRed());
      return;
    }
    if (grand.left === parent) {
      if (parent.right === node) {
        this.rotateL(parent);
        [parent, node] = [node, parent];
      }
      this.rotateR(grand);
    } else {
      if (parent.left === node) {
        this.rotateR(parent);
        [parent, node] = [node, parent];
      }
      this.rotateL(grand);
    }
    parent.toBlack();
    node.toRed();
    grand.toRed();
    return;
  }
  remove(v: T): void {
    let node = this.findNode(this.root, v);
    if (node === null) return;
    if (this.size === 1) {
      this._root = null;
      this._size = 0;
      return;
    }
    if (node.childSize === 2) {
      const successor = this.successor(node)!;
      [node.val, successor.val] = [successor.val, node.val];
      node = successor;
    }
    if (node.childSize === 0) {
      const parent = node.parent!;
      const pos = parent.left === node ? 'left' : 'right';
      parent[pos] = null;
      this._size--;
      this.afterRemove(node);
      return;
    }
    const child = node.left ?? node.right!;
    const parent = node.parent!;
    if (parent === null) {
      this._root = child;
      child.parent = null;
      this._size--;
      this.afterRemove(node);
      return;
    }
    const pos = parent.left === node ? 'left' : 'right';
    parent[pos] = child;
    child.parent = parent;
    this.afterRemove(child);
  }
  afterRemove(node: RBTreeNode2<T>): void {
    if (node.isRed) {
      node.toBlack();
      return;
    }
    let parent = node.parent;
    if (parent === null) return;
    const leftChild = parent.left === null || parent.left === node;
    let sibling = leftChild ? parent.right! : parent.left!;
    if (leftChild) {
      if (sibling.isRed) {
        this.rotateL(parent);
        sibling.toBlack();
        parent.toRed();
        sibling = parent.right!;
      }
      if (sibling.childSize === 0) {
        const parentIsBlack = parent.isBlack;
        sibling.toRed();
        parent.toBlack();
        parentIsBlack && this.afterRemove(parent);
      } else {
        if (isBlack(sibling.right)) {
          this.rotateR(sibling);
          sibling = sibling.parent!;
        }
        this.rotateL(parent);
        sibling.toColor(parent.color);
        sibling.right?.toBlack();
        parent.toBlack();
      }
    } else {
      if (sibling.isRed) {
        this.rotateR(parent);
        sibling.toBlack();
        parent.toRed();
        sibling = parent.left!;
      }
      if (isBlack(sibling.left) && isBlack(sibling.right)) {
        const parentIsBlack = parent.isBlack;
        sibling.toRed();
        parent.toBlack();
        parentIsBlack && this.afterRemove(parent);
      } else {
        if (isBlack(sibling.left)) {
          this.rotateL(sibling);
          sibling = sibling.parent!;
        }
        this.rotateR(parent);
        sibling.toColor(parent.color);
        sibling.left?.toBlack();
        parent.toBlack();
      }
    }
  }
  private successor(node: RBTreeNode2<T>): RBTreeNode2<T> | null {
    if (node.right !== null) {
      let successor = node.right;
      while (successor.left) successor = successor.left;
      return successor;
    }
    let successor = node;
    while (successor.parent !== null && successor.parent.right === successor)
      successor = successor.parent;
    return successor;
  }
  private rotateL(grand: RBTreeNode2<T>) {
    const root = grand.parent;
    const parent = grand.right!;
    if (root === null) this._root = parent;
    else {
      const pos = root.left === grand ? 'left' : 'right';
      root[pos] = parent;
    }
    grand.right = parent.left;
    if (parent.left) parent.left.parent = grand;
    parent.left = grand;
    grand.parent = parent;
    parent.parent = root;
  }
  private rotateR(grand: RBTreeNode2<T>) {
    const root = grand.parent;
    const parent = grand.left!;
    if (root === null) this._root = parent;
    else {
      const pos = root.left === grand ? 'left' : 'right';
      root[pos] = parent;
    }
    grand.left = parent.right;
    if (parent.right) parent.right.parent = grand;
    parent.right = grand;
    grand.parent = parent;
    parent.parent = root;
  }
  print(): string {
    return this.root === null ? '' : this._print(this.root);
  }
  private _print(node: RBTreeNode2<T>, prefix = ''): string {
    const left = node.left;
    const right = node.right;
    const nodeStr = node.toString();
    const halfLength = (nodeStr.length + (prefix === '' ? 0 : 3)) >> 1;
    const blankStr = repeat(' ', halfLength);
    const lineStr = repeat('─', halfLength);
    prefix += blankStr;
    let str = nodeStr + `\n`;
    if (right !== null) {
      str +=
        prefix +
        (left === null ? '└' : '├') +
        lineStr +
        ' R ' +
        this._print(right, `${prefix}${left === null ? ' ' : '│'}${blankStr}`);
    }
    if (left !== null) {
      str += `${prefix}└${lineStr}` + ' L ' + this._print(left, `${prefix}${blankStr} `);
    }
    return str;
  }
  levelOrder(): T[] {
    const list: T[] = [];
    if (this.root === null) return list;
    const queue: RBTreeNode2<T>[] = [this.root];
    while (queue.length !== 0) {
      const node = queue.shift()!;
      list.push(node.val);
      node.left && queue.push(node.left);
      node.right && queue.push(node.right);
    }
    return list;
  }
}

核心代码​

/* istanbul ignore file */
import { repeat } from 'lodash';
import { IRBTree } from './rbTree';
const RED = true;
const BLACK = false;
type Color = boolean;
const isRed = (node: RBTreeNode3<any> | null) => node?.color === RED;
const isBlack = (node: RBTreeNode3<any> | null) => node === null || node?.color === BLACK;
const setColor = (node: RBTreeNode3<any> | null, color: Color) => {
  if (node) node.color = color;
};
const setRed = (node: RBTreeNode3<any> | null) => setColor(node, RED);
const setBlack = (node: RBTreeNode3<any> | null) => setColor(node, BLACK);
export class RBTreeNode3<T> {
  color = RED;
  left: RBTreeNode3<T> | null = null;
  right: RBTreeNode3<T> | null = null;
  constructor(public val: T, public parent: RBTreeNode3<any> | null = null) {}
  toString() {
    return this.val + (isRed(this) ? '(R)' : '');
  }
}
export class RBTree3<T> implements IRBTree<T> {
  private _size = 0;
  get size() {
    return this._size;
  }
  get empty() {
    return this._size === 0;
  }
  private root: RBTreeNode3<T> | null = null;
  constructor(private compare: (k1: T, k2: T) => number) {}
  clear() {
    this.root = null;
    this._size = 0;
  }
  contains(key: T) {
    return this.findNode(key) !== null;
  }
  add(key: T) {
    let node = this.findNode(key);
    if (this.root === null) {
      this.root = node = new RBTreeNode3<T>(key);
    } else {
      let parent = this.root;
      let pos = 'left';
      while (parent !== null) {
        const compare = this.compare(parent.val, key);
        if (compare > 0) {
          if (parent.left === null) break;
          parent = parent.left;
        } else {
          if (parent.right === null) {
            pos = 'right';
            break;
          }
          parent = parent.right;
        }
      }
      parent[pos] = node = new RBTreeNode3<T>(key, parent);
    }
    this._size++;
    this.afterSet(node);
  }
  private afterSet(node: RBTreeNode3<T>) {
    let parent = node.parent;
    if (parent === null) {
      setBlack(node);
      return;
    }
    if (isBlack(parent)) return;
    const grand = parent.parent!;
    const sibling = grand.left === parent ? grand.right! : grand.left!;
    if (isRed(sibling)) {
      setBlack(parent);
      setBlack(sibling);
      setRed(grand);
      this.afterSet(grand);
      return;
    }
    if (grand.left === parent) {
      if (parent.right === node) {
        this.rotateL(parent);
        [parent, node] = [node, parent];
      }
      this.rotateR(grand);
    } else {
      if (parent.left === node) {
        this.rotateR(parent);
        [parent, node] = [node, parent];
      }
      this.rotateL(grand);
    }
    setBlack(parent);
    setRed(node);
    setRed(grand);
  }
  remove(key: T) {
    let node = this.findNode(key);
    if (node === null) return;
    if (node.left !== null && node.right !== null) {
      const successor = this.successor(node);
      [node.val, successor.val] = [successor.val, node.val];
      node = successor;
    }
    const parent = node.parent!;
    this._size--;
    if (node.left === null && node.right === null) {
      if (this.root === node) {
        this.clear();
        return;
      }
      const pos = parent.left === node ? 'left' : 'right';
      parent[pos] = null;
      this.afterRemove(node);
      return;
    }
    const childNode = node.left ?? node.right!;
    if (parent === null) this.root = childNode;
    else {
      const pos = parent.left === node ? 'left' : 'right';
      parent[pos] = childNode;
    }
    childNode.parent = parent;
    this.afterRemove(childNode);
  }
  private afterRemove(node: RBTreeNode3<T>) {
    if (isRed(node)) {
      setBlack(node);
      return;
    }
    const parent = node.parent;
    if (parent === null) return;
    const leftChild = parent.left === null || parent.left === node;
    let sibling = leftChild ? parent.right! : parent.left!;
    if (leftChild) {
      if (isRed(sibling)) {
        this.rotateL(parent);
        setBlack(sibling);
        setRed(parent);
        sibling = parent.right!;
      }
      if (isBlack(sibling.left) && isBlack(sibling.right)) {
        const parentIsBlack = isBlack(parent);
        setRed(sibling);
        setBlack(parent);
        parentIsBlack && this.afterRemove(parent);
        return;
      }
      if (isBlack(sibling.right)) {
        this.rotateR(sibling);
        sibling = sibling.parent!;
      }
      this.rotateL(parent);
      setColor(sibling, parent.color);
      setBlack(parent);
      setBlack(sibling.right);
    } else {
      if (isRed(sibling)) {
        this.rotateR(parent);
        setBlack(sibling);
        setRed(parent);
        sibling = parent.left!;
      }
      if (isBlack(sibling.left) && isBlack(sibling.right)) {
        const parentIsBlack = isBlack(parent);
        setRed(sibling);
        setBlack(parent);
        parentIsBlack && this.afterRemove(parent);
        return;
      }
      if (isBlack(sibling.left)) {
        this.rotateL(sibling);
        sibling = sibling.parent!;
      }
      this.rotateR(parent);
      setColor(sibling, parent.color);
      setBlack(parent);
      setBlack(sibling.left);
    }
  }
  private successor(node: RBTreeNode3<T>) {
    let successor = node.right!;
    while (successor.left) successor = successor.left;
    return successor;
  }
  private findNode(key: T, root = this.root): RBTreeNode3<T> | null {
    if (root === null) return null;
    const compare = this.compare(root.val, key);
    if (compare > 0) return this.findNode(key, root.left);
    if (compare < 0) return this.findNode(key, root.right);
    return root;
  }
  private rotateL(grand: RBTreeNode3<T>) {
    const root = grand.parent;
    const parent = grand.right!;
    if (root === null) this.root = parent;
    else {
      const pos = root.left === grand ? 'left' : 'right';
      root[pos] = parent;
    }
    grand.right = parent.left;
    if (parent.left) parent.left.parent = grand;
    parent.left = grand;
    grand.parent = parent;
    parent.parent = root;
  }
  private rotateR(grand: RBTreeNode3<T>) {
    const root = grand.parent;
    const parent = grand.left!;
    if (root === null) this.root = parent;
    else {
      const pos = root.left === grand ? 'left' : 'right';
      root[pos] = parent;
    }
    grand.left = parent.right;
    if (parent.right) parent.right.parent = grand;
    parent.right = grand;
    grand.parent = parent;
    parent.parent = root;
  }
  print(): string {
    return this.root === null ? '' : this._print(this.root);
  }
  private _print(node: RBTreeNode3<T>, prefix = ''): string {
    const left = node.left;
    const right = node.right;
    const nodeStr = node.toString();
    const halfLength = (nodeStr.length + (prefix === '' ? 0 : 3)) >> 1;
    const blankStr = repeat(' ', halfLength);
    const lineStr = repeat('─', halfLength);
    prefix += blankStr;
    let str = nodeStr + `\n`;
    if (right !== null) {
      str +=
        prefix +
        (left === null ? '└' : '├') +
        lineStr +
        ' R ' +
        this._print(right, `${prefix}${left === null ? ' ' : '│'}${blankStr}`);
    }
    if (left !== null) {
      str += `${prefix}└${lineStr}` + ' L ' + this._print(left, `${prefix}${blankStr} `);
    }
    return str;
  }
  levelOrder(): T[] {
    const list: T[] = [];
    if (this.root === null) return list;
    const queue: RBTreeNode3<T>[] = [this.root];
    while (queue.length !== 0) {
      const node = queue.shift()!;
      list.push(node.val);
      node.left && queue.push(node.left);
      node.right && queue.push(node.right);
    }
    return list;
  }
}