web/rbt.js

const RED = 0;
const BLACK = 1;

function isDoubleBlack(x) {
    return x && nodeDifferences(x).indexOf(2) != -1;
}

class RBTree extends RankedTree {
    isCorrectNode(node, recursive) {
        if (!node) {
            return true;
        }

        let [left, right] = nodeDifferences(node);
        if ([RED, BLACK].indexOf(left) == -1) {
            // left subtree has invalid difference
            return false;
        } else if ([RED, BLACK].indexOf(right) == -1) {
            // right subtree has invalid difference
            return false;
        }

        if (nodeDifference(node) == RED && (left == RED || right == RED)) {
            // two consecutive red nodes
            return false;
        }

        recursive = recursive ?? true;
        return (
            !recursive ||
            (this.isCorrectNode(node.left) && this.isCorrectNode(node.right))
        );
    }

    insertRebalanceStep(z, y, rightChild, rotateLeft, rotateRight) {
        let [p, pp] = [z.parent, z.parent.parent];

        if (y && nodeDifference(y) == RED) {
            // Case 1
            // ======
            // z’s uncle y is red
            pp.rank++;
            this.record("Insertion case #1: recoloring uncle", pp);
            z = pp;
        } else if (z === rightChild) {
            // Case 2
            // ======
            // z’s uncle y is black and z is a right child
            z = p;
            rotateLeft(this, p);
            this.record("Insertion case #2: rotating by parent");
        } else {
            // Case 3
            // ======
            // z’s uncle y is black and z is a left child
            rotateRight(this, pp);
            this.record("Insertion case #3: rotating by grandparent");
        }

        return z;
    }

    insertRebalance(z) {
        while (z.parent && nodeDifference(z.parent) == RED) {
            let [p, pp] = [z.parent, z.parent.parent];

            if (p === pp.left) {
                z = this.insertRebalanceStep(
                    z,
                    pp.right,
                    p.right,
                    rotateLeft,
                    rotateRight
                );
            } else {
                z = this.insertRebalanceStep(
                    z,
                    pp.left,
                    p.left,
                    rotateRight,
                    rotateLeft
                );
            }
        }
    }

    deleteRebalanceStep(x, w, parent, right, rotateLeft, rotateRight) {
        if (nodeDifference(w) == RED) {
            // Case 1
            // ======
            // x’s sibling w is red
            rotateLeft(this, parent);
            this.record("Deletion case #1: rotating by parent", parent);
            w = right(parent);
        }

        if (nodeDifferences(w).equals([BLACK, BLACK])) {
            // Case 2
            // ======
            // x’s sibling w is black, and both of w’s children are black
            parent.rank--;
            this.record("Deletion case #2: recoloring sibling", w);
            x = parent;
        } else {
            // Case 3
            // ======
            // x’s sibling w is black,
            // w’s left child is red, and w’s right child is black
            if (nodeDifference(right(w), w) == BLACK) {
                rotateRight(this, w);
                this.record("Deletion case #3: rotating by w", w);
                w = right(parent);
            }

            // Case 4
            // ======
            // x’s sibling w is black, and w’s right child is red
            parent.rank--;
            this.record(
                "Deletion case #4: moving double black to parent",
                parent
            );
            w.rank++;
            this.record("Deletion case #4: coloring w's child", w);
            rotateLeft(this, parent);
            this.record("Deletion case #4: rotating by parent", parent);

            x = this.root;
        }

        return [x, x ? x.parent : null];
    }

    deleteRebalance(node, parent) {
        if (!node && !parent) {
            return;
        }

        while (node !== this.root && isDoubleBlack(parent)) {
            if (node === parent.left) {
                [node, parent] = this.deleteRebalanceStep(
                    node,
                    parent.right,
                    parent,
                    (x) => x.right,
                    rotateLeft,
                    rotateRight
                );
            } else {
                [node, parent] = this.deleteRebalanceStep(
                    node,
                    parent.left,
                    parent,
                    (x) => x.left,
                    rotateRight,
                    rotateLeft
                );
            }
        }
    }

    styleEdge(rankDifference) {
        switch (rankDifference) {
            case 2:
                return " penwidth=2";
            case 0:
                return ', color="red"';
            case 1:
            default:
                return "";
        }
    }
}