红黑树js

let root = null;
function RBTNode(key, parent) {
  this.key = key;
  this.parent = parent;
  this.left = null;
  this.right = null;
  this.color = true;//默认插入的节点颜色为红色
}

// 辅助函数
// 使用boolean值作为颜色常量
const RED = false;
const BLACK = true;

//染色
const color = (node, color) => {
  if(node == null) return node;
  node.color = color;
  return node;
};
//将节点染红色
const red = (node) => color(node, RED);
//将节点染黑色
const black = (node) => color(node, BLACK);

//获取节点颜色
const colorOf = node => {
  return node == null ? BLACK : node.color;
}
//判断节点颜色
const isBlack = (node) => colorOf(node) == BLACK;
const isRed = (node) => colorOf(node) == RED;
//获取左节点
const left = node => node == null ? null : node.left;
const right = node => node == null ? null : node.right;

//判断节点在父节点的左子树上
const isLeftChild = (node) => {
  return node && node.parent && node.parent.left == node;
};
//判断节点在父节点的右子树上
const isRightChild = (node) => node && node.parent && node.parent.right == node;
//判断是否度为2的节点
const hasTwoChildren = node => node && node.left != null && node.right != null;

//获取兄弟节点
const sibling = node => {
  if(isLeftChild(node)){
    return node.parent.right
  }
  if(isRightChild(node)) {
    return node.parent.left
  }
  return null;
};

//递归查找值为key的节点
const find = (node, key) => {
  if(node == null) {
    return null;
  }
  if(node.key == key) {
    return node;
  }
  return find(key > node.key ? node.right : node.left, key);
}

//查找值为key的节点
const search = (key) => find(tRoot, key);
/**
 * 添加节点
 */
const add = (key, value) => {
  //添加
  if(root == null) {
    root = new RBTNode(key, null);
    afterAdd(root);
    return;
  }
  let parent = null;
  let cur = root;
  while (cur != null) {
    if(key > cur.key) {//右子树上
      parent = cur;
      cur = cur.right
    }else if(key < cur.key) {//左子树上
      parent = cur;
      cur = cur.left
    }else  {
    //  key == cur.key 不处理或更新节点
    }
  }
  
  //找到底层节点了
  let newNode = new RBTNode(key, parent);
  if(key > parent.key) {
    parent.right = newNode
  }else if (key < parent.key){
    parent.left = newNode;
  }
  
  //添加后 维持红黑树的性质:染色和旋转
  afterAdd(newNode);
  return;
}
/**
 * 添加节点后维护红黑树性质
 * @param node:新节点
 */
const afterAdd = node => {
  let parent = node.parent;
  //根节点只需染黑色
  if(parent == null) {
    black(node);
    return;
  }
  //看父节点是黑色还是红色?
  //父节点为黑色,无需处理
  if(isBlack(parent)) {
    return;
  }
  
  //父节点是红色
  //就要看叔父节点
  let uncle = sibling(parent);
  let grand = parent.parent;
  
  //看叔父节点是红色还是黑色?
  //叔父节点红色: 父节点和叔父节点都染黑色,祖父节点染红色,再将祖父节点当作新增节点,继续红黑树的调整
  if(isRed(uncle)) {
    black(uncle);
    black(parent);
    red(grand);
    afterAdd(grand);
    return;
  }
  
  //叔父节点不是红色(是黑色的)
  //看父节点是祖父节点的左子节点还是右子节点?
  if(isLeftChild(parent)) {//L
    if(isLeftChild(node)) {//LL
      black(parent);
      red(grand);
      //grand进行LL旋转
      LLrotate(grand);
    }else {//LR
      black(node);
      red(grand);
      LRrotate(grand);
    }    
  }else {//R
    if(isRightChild(node)) {//RR
      black(parent);
      red(grand);
      //grand进行RR旋转
      RRrotate(grand);      
    }else {//RL
      black(node);
      red(grand);
      RLrotate(grand);    
    }    
  }   
}

/**
* @description:红黑树的左旋RR
* @param {type} 插入的叶子节点
* @return: 返回旋转后的根节点
*/
const RRrotate = node => {
  let newRoot = node.right;
  node.right = newRoot.left;
  if(newRoot.left != null) {
    newRoot.left.parent = node;
  } 
  
  newRoot.left = node;
  newRoot.parent = node.parent;
  
  if(isLeftChild(node)) {
    node.parent.left = newRoot;
  }else if(isRightChild(node)) {
    node.parent.right = newRoot;
  }else {
    //node是根节点
    root = newRoot;
  }
  
  node.parent = newRoot;  
  
  return newRoot;
}
/**
 * 红黑树的右旋LL
 */
const LLrotate = node => {
  let newRoot = node.left;
  node.left = newRoot.right;
  if(newRoot.right != null) {
    newRoot.right.parent = node;
  }
  newRoot.right = node;
  newRoot.parent = node.parent;
  if(isLeftChild(node)) {
    node.parent.left = newRoot;
  }else if(isRightChild(node)) {
    node.parent.right = newRoot;
  }else {
    root = newRoot;
  }  
  
  node.parent = newRoot;
  return newRoot;
}

const RLrotate = node => {
  LLrotate(node.right);
  return RRrotate(node);
};

const LRrotate = node => {
  RRrotate(node.left);
  return LLrotate(node);
};

/***
 * 删除节点
 */
const remove = key => {
  let node = node(key);
  removeNode(node);
};

const removeNode = node => {
  if(node == null) {
    return;
  }
  
  //度为2的节点
  if(hasTwoChildren(node)) {
    //使用前驱 或 后继的值去替换节点的值,实际删除前驱或后继
    //查找前驱 node.left.right.right...
    let p = node.left;
    while (p.right != null) {
      p = p.right
    }
    node.key = p.key;//将p的数据置换到node上
    node = p;//将node指针指向p,后面删除node
  }
  
  //删除度为0和度为1的
  
  //找到替代节点,度为0,replacement = null;
  let replacement = node.left != null ? node.left : node.right;
  
  if(node.parent == null) {
    //node为根节点,即删除的是根节点
    root = replacement;//replacement此时肯定为null
  }else if(isLeftChild(node)) {
    node.parent.left = replacement;
  }else if(isRightChild(node)) {
    node.parent.right = replacement;
  }
  
  if(replacement != null) {
    replacement.parent = node.parent;
  }
  //删除节点后维持红黑树的性质: 染色和旋转
  afterRemove(node, replacement);
  node.left = null;
  node.right = null;
  node.parent = null;  
};

const afterRemove = (node, replacement) => {
  //1.删除节点是红色,无需处理
  if(isRed(node)) {
    return;
  }
  //2.删除节点为黑色
  //2-1.度为1, 有红色替换子节点,将子节点染黑色即可
  if(isRed(replacement)){
    black(replacement);
    return;
  }
  
  //2-2.度为0,
  //2-2-1.如果删除的是根节点,无需处理
  let parent = node.parent;
  if(parent == null) {
    return;
  }
  
  //要删除的是黑色的非根叶子节点
  //删除节点是在左边还是右边??
  //节点在parent左子树上,删除后,左子树指向null,以此判断节点是否在左子树上
  let left = parent.left == null || isLeftChild(node);
  let sibling = left ? parent.right: parent.left;
  
  if(left) {
    //兄弟节点是红色, 兄弟染黑色,父节点染红色,在进行RR旋转, 换兄弟
    if(isRed(sibling)) {
      black(sibling);
      red(parent);
      RRrotate(parent);
      //换兄弟
      sibling = parent.right;
    }
    
    //兄弟为黑色
    //兄弟没有红色子节点
    if(isBlack(sibling.left) && isBlack(sibling.right)) {
      //父节点的颜色
      let parentBlack = isBlack(parent);
      red(sibling);
      black(parent);
      if(parentBlack) {
        //父节点原本是黑色的,将父节点当删除节点继续向上维护红黑树性质
        afterRemove(parent, null);
      }
    }else {
      let newParent;
      if(isRed(sibling.right)) {//RR
        newParent = RRrotate(parent)
      }else {//RL
        newParent = RLrotate(parent);
      }
      
      color(newParent, colorOf(parent));
      black(newParent.left);
      black(newParent.right);      
    }
    
  }else {
     if(isRed(sibling)) {
      black(sibling);
      red(parent);
      LLrotate(parent);
      //换兄弟, 
      sibling = parent.left;
    }
    
    //兄弟为黑色
    //兄弟没有红色子节点
    if(isBlack(sibling.left) && isBlack(sibling.right)) {
      //父节点的颜色
      let parentBlack = isBlack(parent);
      red(sibling);
      black(parent);
      if(parentBlack) {
        //父节点原本是黑色的,将父节点当删除节点继续向上维护红黑树性质
        afterRemove(parent, null);
      }
    }else {
      let newParent;
      if(isRed(sibling.left)) {//LL
        newParent = LLrotate(parent)
      }else {//LR
        newParent = LRrotate(parent);
      }
      //新的父节点继承原来父节点的颜色,新的父节点的左右子节点染黑色
      color(newParent, colorOf(parent));
      black(newParent.left);
      black(newParent.right);      
    }
  }
  
}

/**
 * 查找被删除的节点
 */
const node = key => {
  let cur = root;
  while (cur != null) {
    if(key > cur.key) {
      cur = cur.right;
    }else if(key < cur.key) {
      cur = cur.left;
    }else {
      //cur.key==key,找到了
      return cur;
    }
  }
  return null;
}