sonic-forest/lib/red-black/util.js

High-performance (binary) tree and sorted map implementation (AVL, Splay, Radix, Red-Black)

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.remove = exports.insertLeft = exports.insertRight = exports.insert = exports.print = void 0;
const swap_1 = require("../util/swap");
const print_1 = require("../util/print");
Object.defineProperty(exports, "print", { enumerable: true, get: function () { return print_1.print; } });
const insert = (root, n, comparator) => {
    if (!root)
        return (n.b = true), n;
    const key = n.k;
    let curr = root;
    let next = undefined;
    let cmp = 0;
    while ((next = ((cmp = comparator(key, curr.k)) < 0 ? curr.l : curr.r)))
        curr = next;
    return (cmp < 0 ? (0, exports.insertLeft)(root, n, curr) : (0, exports.insertRight)(root, n, curr));
};
exports.insert = insert;
const insertRight = (root, n, p) => {
    const g = p.p;
    p.r = n;
    n.p = p;
    if (p.b || !g)
        return root;
    const top = rRebalance(n, p, g);
    return top.p ? root : top;
};
exports.insertRight = insertRight;
const insertLeft = (root, n, p) => {
    const g = p.p;
    p.l = n;
    n.p = p;
    if (p.b || !g)
        return root;
    const top = lRebalance(n, p, g);
    return top.p ? root : top;
};
exports.insertLeft = insertLeft;
const rRebalance = (n, p, g) => {
    const gl = g.l;
    const zigzag = gl === p;
    const u = zigzag ? g.r : gl;
    const uncleIsBlack = !u || u.b;
    if (uncleIsBlack) {
        g.b = false;
        if (zigzag) {
            lrRotate(g, p, n);
            n.b = true;
            return n;
        }
        p.b = true;
        rRotate(g, p);
        return p;
    }
    return recolor(p, g, u);
};
const lRebalance = (n, p, g) => {
    const gr = g.r;
    const zigzag = gr === p;
    const u = zigzag ? g.l : gr;
    const uncleIsBlack = !u || u.b;
    if (uncleIsBlack) {
        g.b = false;
        if (zigzag) {
            rlRotate(g, p, n);
            n.b = true;
            return n;
        }
        p.b = true;
        lRotate(g, p);
        return p;
    }
    return recolor(p, g, u);
};
const recolor = (p, g, u) => {
    p.b = true;
    if (u)
        u.b = true;
    const gg = g.p;
    if (gg) {
        g.b = false;
        if (gg.b)
            return g;
    }
    else {
        g.b = true;
        return g;
    }
    const ggg = gg.p;
    if (!ggg)
        return gg;
    return gg.l === g ? lRebalance(g, gg, ggg) : rRebalance(g, gg, ggg);
};
const lRotate = (n, nl) => {
    const p = n.p;
    const nlr = nl.r;
    ((nl.r = n).l = nlr) && (nlr.p = n);
    ((n.p = nl).p = p) && (p.l === n ? (p.l = nl) : (p.r = nl));
};
const rRotate = (n, nr) => {
    const p = n.p;
    const nrl = nr.l;
    ((nr.l = n).r = nrl) && (nrl.p = n);
    ((n.p = nr).p = p) && (p.l === n ? (p.l = nr) : (p.r = nr));
};
const lrRotate = (g, p, n) => {
    const gg = g.p;
    const nl = n.l;
    const nr = n.r;
    gg && (gg.l === g ? (gg.l = n) : (gg.r = n));
    n.p = gg;
    n.l = p;
    n.r = g;
    p.p = g.p = n;
    (p.r = nl) && (nl.p = p);
    (g.l = nr) && (nr.p = g);
};
const rlRotate = (g, p, n) => {
    const gg = g.p;
    const nl = n.l;
    const nr = n.r;
    gg && (gg.l === g ? (gg.l = n) : (gg.r = n));
    n.p = gg;
    n.l = g;
    n.r = p;
    g.p = p.p = n;
    (g.r = nl) && (nl.p = g);
    (p.l = nr) && (nr.p = p);
};
const remove = (root, n) => {
    const originalNode = n;
    const r = n.r;
    const l = n.l;
    let child;
    if (r) {
        let inOrderSuccessor = r;
        if (inOrderSuccessor)
            while (true) {
                const next = inOrderSuccessor.l;
                if (next)
                    inOrderSuccessor = next;
                else
                    break;
            }
        n = inOrderSuccessor;
        child = n.r;
    }
    else if (!n.p) {
        if (l) {
            l.b = true;
            l.p = undefined;
        }
        return l;
    }
    else {
        child = r || l;
    }
    if (n !== originalNode) {
        originalNode.k = n.k;
        originalNode.v = n.v;
        const b = n.b;
        n.b = originalNode.b;
        originalNode.b = b;
        root = (0, swap_1.swap)(root, originalNode, n);
        n = originalNode;
    }
    if (child) {
        const p = n.p;
        child.p = p;
        if (p.l === n)
            p.l = child;
        else
            p.r = child;
        if (!child.b)
            child.b = true;
        else
            root = correctDoubleBlack(root, child);
    }
    else {
        if (n.b)
            root = correctDoubleBlack(root, n);
        const p2 = n.p;
        if (p2) {
            if (n === p2.l)
                p2.l = undefined;
            else
                p2.r = undefined;
        }
        else {
            n.b = true;
            return n;
        }
    }
    return root;
};
exports.remove = remove;
const correctDoubleBlack = (root, n) => {
    LOOP: while (true) {
        const p = n.p;
        if (!p)
            return n;
        const pl = p.l;
        const isLeftChild = n === pl;
        let s = (isLeftChild ? p.r : pl);
        const sl = s.l;
        if (s && !s.b && (!sl || sl.b)) {
            const sr = s.r;
            if (!sr || sr.b) {
                if (isLeftChild)
                    rRotate(p, s);
                else
                    lRotate(p, s);
                p.b = false;
                s.b = true;
                if (!s.p)
                    root = s;
            }
        }
        if (p.b && s.b && (!sl || sl.b)) {
            const sr = s.r;
            if (!sr || sr.b) {
                s.b = false;
                n = p;
                continue LOOP;
            }
        }
        if (!p.b) {
            const pl = p.l;
            s = (n === pl ? p.r : pl);
            const sl = s.l;
            if (s.b && (!sl || sl.b)) {
                const sr = s.r;
                if (!sr || sr.b) {
                    const sr = s.r;
                    if (!sr || sr.b) {
                        s.b = false;
                        p.b = true;
                        return root;
                    }
                }
            }
        }
        if (s.b) {
            const sl = s.l;
            const sr = s.r;
            if (n === p.l && (!sr || sr.b) && sl && !sl.b) {
                sl.b = true;
                s.b = false;
                lRotate(s, sl);
            }
            else if (n === p.r && (!sl || sl.b) && sr && !sr.b) {
                sr.b = true;
                s.b = false;
                rRotate(s, sr);
            }
            if (!s.p)
                return s;
            const pl = p.l;
            s = (n === pl ? p.r : pl);
        }
        s.b = p.b;
        p.b = true;
        if (n === p.l) {
            s.r.b = true;
            rRotate(p, s);
        }
        else {
            s.l.b = true;
            lRotate(p, s);
        }
        return s.p ? root : s;
    }
};