RBTree.js

const RED = true;
const BLACK = false;

// 定义红黑树节点
class Node {
    constructor(key, value) {
        this.key = key; // 键
        this.value = value; // 值
        this.color = RED; // 颜色
        this.left = null; // 左子节点
        this.right = null; // 右子节点
    }
}

module.exports = class RBTree {
    constructor() {
        this.root = null;
        this.size = 0;
    }
    // 判断节点node的颜色
    isRed(node) {
        if(node === null)
            return BLACK;
        return node.color;
    }

    //      node                            x
    //     /    \         左旋转           /  \   
    //    T1     x   ---------------->   node  T3
    //          /  \                    /   \
    //         T2  T3                  T1   T2
    leftRotate(node) {
        let x = node.right;
        // 左旋转
        node.right = x.left;
        x.left = node;
        // 更新颜色;
        x.color = node.color;
        node.color = RED;

        return x;
    }

    //      node                            x
    //     /    \         右旋转           /  \   
    //    x     T2   ---------------->   y    node
    //   /  \                                /   \
    //  y   T1                              T1    T2
    rightRotate(node) {
        let x = node.left;
        // 右旋转;
        node.left = x.right;
        x.right = node;
        // 更新颜色;
        x.color = node.color;
        node.color = RED;

        return x;
    }
    // 颜色翻转
    flipColors(node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
    }
    // 向二分搜索树中添加新的元素(key, value)
    add(key, value) {
        this.root = this._add(this.root, key, value);
        this.root.color = BLACK; // 最终根节点为黑色节点
    }
    // 向以node为根的二分搜索树中插入元素(key, value),递归算法
    // 返回插入新节点后红黑树的根
    _add(node, key, value) {
        if(node === null) {
            this.size++;
            return new Node(key, value); // 默认插入红色节点
        }

        if(key < node.key) {
            node.left = this._add(node.left, key, value);
        } else if(key > node.key) {
            node.right = this._add(node.right, key, value);
        } else {
            node.value = value; // 更新值
        }
        // 利用回溯法更新颜色
        // 颜色变化情况： B -- 黑色  R -- 红色  n -- null(叶子节点)
        // 1. 插入'A'
        //     A(B)
        //    /    \
        // n(B)    n(B)
        // 
        // 2. 插入'F': 需要变换，使用方法1变换
        //     A(B)                       F(B)
        //    /    \            左旋      /   \
        //  n(B)   F(R)       ------->  A(R)   n(B)
        //        /    \               /   \
        //       n(B)  n(B)           n(B)  n(B)
        // 
        // 3. 插入 'M'
        //         F(B)          
        //       /     \          
        //     A(R)     M(R)      
        //    /   \     /   \      
        // n(B)  n(B)  n(B)  n(B)  
        // 
        // 4. 插入 'C': 需要变换，使用方法1变换,之后再用方法2变换,最后，再用方法3翻转颜色,最后根节点变换为黑色
        //         F(B)                            F(B)                               C(B)                                 C(R)
        //       /     \                         /      \                           /      \                             /      \
        //     A(R)     M(R)       左旋         C(R)      M(R)           右旋       A(R)    F(R)           翻转颜色      A(B)     F(B)
        //    /   \     /   \     ------->    /   \      /    \       -------->  /   \     /    \         --------->  /   \     /    \
        // n(B)  C(R)  n(B)  n(B)           A(R)  n(B)  n(B)   n(B)            n(B) n(B)  n(B)  M(R)                n(B)  n(B) n(B)  M(R)
        //      /    \                     /   \                                               /   \                                /   \
        //    n(B)   n(B)                n(B)  n(B)                                           n(B)  n(B)                          n(B)  n(B) 

        // 方法1：如果当前节点的右节点为红色，且左节点非红色，则进行左旋
        if(this.isRed(node.right) && !this.isRed(node.left)) {
            node = this.leftRotate(node);
        }
        // 方法2：如果当前节点的左节点和左左节点都为红色，则进行右旋
        if(this.isRed(node.left) && this.isRed(node.left.left)) {
            node = this.rightRotate(node);
        }
        // 方法3：如果当前节点左右子树都为红色，则翻转颜色
        if(this.isRed(node.left) && this.isRed(node.right))
            this.flipColors(node);
        return node;
    }
    // 返回以node为根节点的二分搜索树中，key所在的节点
    getNode(node, key) {
        if(node === null) {
            return null;
        }
        if(key === node.key) {
            return node;
        } else if(key < node.key) {
            return this.getNode(node.left, key);
        } else {
            return this.getNode(node.right, key);
        }
    }
    // 判断是包含该值
    contains(key) {
        return this.getNode(this.root, key) !== null;
    }
    get(key) {
        let node = this.getNode(this.root, key);
        return node === null ? null : node.value;
    }
    set(key, value) {
        let node = this.getNode(this.root, key);
        if(node === null) {
            throw(key+" doesn't exist!");
        }
        node.value = value;
    }
    // 返回以node为根的二分搜索树的最小值所在的节点
    minimun(node) {
        if(node.left === null)
            return node;
        return this.minimun(node.left);
    }
    // 删除掉以node为根的二分搜索树中的最小节点
    // 返回删除节点后新的二分搜索树的根
    removeMin(node) {
        if(node.left === null) {
            let rightNode = node.right;
            node.right = null;
            this.size--;
            return rightNode;
        }
        node.left = this.removeMin(node.left);
        return node;
    }
    // 从二分搜索树中删除键为key的节点
    remove(key) {
        let node = this.getNode(this.root, key);
        if(node !== null) {
            this.root = this._remove(this.root, key);
            return node.value;
        }
        return null;
    }
    _remove(node, key) {
        if(node === null) {
            return null;
        }
        if(key < node.key) {
            node.left = this._remove(node.left, key);
            return node;
        }
        else if(key > node.key) {
            node.right = this._remove(node.right, key);
            return node;
        }
        else {
            // 待删除节点左子树为空的情况
            if(node.left === null) {
                let rightNode = node.right;
                node.right = null;
                this.size--;
                return rightNode;
            }
            // 待删除节点右子树为空的情况
            if(node.right === null) {
                let leftNode = node.left;
                node.left = null;
                this.size--;
                return leftNode;
            }
            // 待删除节点左右子树均不为空的情况
            // 找到比待删除节点大的最小节点，即待删除节点右子树的最小节点
            // 用这个节点顶替待删除节点的位置
            let successor = this.minimun(node.right);
            successor.right = this.removeMin(node.right);
            successor.left = node.left;
            node.left = node.right = null;
            return successor;
        }
    }
    // 返回size
    getSize() {
        return this.size;
    }
    // 判断是否为空
    isEmpty() {
        return this.size === 0;
    }
}