tdigest-sd/lib/rbtree.js

An implementation of tdigest with serialization and deserialization for space-efficient storage.

var TreeBase = require('./treebase');

function Node(data) {
    this.data = data;
    this.left = null;
    this.right = null;
    this.red = true;
}

Node.serialize = function(node, data_serializer) {
    if (node) {
        var d = typeof data_serializer === 'function' ? data_serializer(node.data) : node.data;
        return [
            d, Node.serialize(node.left, data_serializer),
            Node.serialize(node.right, data_serializer), node.red ? 1 : 0
        ];
    } else {
        return 0;
    }
};

Node.deserialize = function(arr, data_deserializer) {
    var node = null;
    if (arr) {
        var d = typeof data_deserializer === 'function' ? data_deserializer(arr[0]) : arr[0];
        node = new Node(d);
        node.left = Node.deserialize(arr[1], data_deserializer);
        node.right = Node.deserialize(arr[2], data_deserializer);
        node.red = arr[3] === 1;
    }
    return node;
};

Node.prototype.get_child = function(dir) {
    return dir ? this.right : this.left;
};

Node.prototype.set_child = function(dir, val) {
    if (dir) {
        this.right = val;
    } else {
        this.left = val;
    }
};

function RBTree(comparator) {
    this._root = null;
    this._comparator = comparator;
    this.size = 0;
}

RBTree.prototype = new TreeBase();

RBTree.prototype.deserialize = function(arr, data_deserializer) {
    this.size = arr[0];
    this._root = Node.deserialize(arr[1], data_deserializer);
};

RBTree.prototype.serialize = function(data_serializer) {
    return [this.size, Node.serialize(this._root, data_serializer)];
};

// returns true if inserted, false if duplicate
RBTree.prototype.insert = function(data) {
    var ret = false;

    if (this._root === null) {
        // empty tree
        this._root = new Node(data);
        ret = true;
        this.size++;
    } else {
        var head = new Node(undefined); // fake tree root

        var dir = 0;
        var last = 0;

        // setup
        var gp = null; // grandparent
        var ggp = head; // grand-grand-parent
        var p = null; // parent
        var node = this._root;
        ggp.right = this._root;

        // search down
        while (true) {
            if (node === null) {
                // insert new node at the bottom
                node = new Node(data);
                p.set_child(dir, node);
                ret = true;
                this.size++;
            } else if (is_red(node.left) && is_red(node.right)) {
                // color flip
                node.red = true;
                node.left.red = false;
                node.right.red = false;
            }

            // fix red violation
            if (is_red(node) && is_red(p)) {
                var dir2 = ggp.right === gp;

                if (node === p.get_child(last)) {
                    ggp.set_child(dir2, single_rotate(gp, !last));
                } else {
                    ggp.set_child(dir2, double_rotate(gp, !last));
                }
            }

            var cmp = this._comparator(node.data, data);

            // stop if found
            if (cmp === 0) {
                break;
            }

            last = dir;
            dir = cmp < 0;

            // update helpers
            if (gp !== null) {
                ggp = gp;
            }
            gp = p;
            p = node;
            node = node.get_child(dir);
        }

        // update root
        this._root = head.right;
    }

    // make root black
    this._root.red = false;

    return ret;
};

// returns true if removed, false if not found
RBTree.prototype.remove = function(data) {
    if (this._root === null) {
        return false;
    }

    var head = new Node(undefined); // fake tree root
    var node = head;
    node.right = this._root;
    var p = null; // parent
    var gp = null; // grand parent
    var found = null; // found item
    var dir = 1;

    while (node.get_child(dir) !== null) {
        var last = dir;

        // update helpers
        gp = p;
        p = node;
        node = node.get_child(dir);

        var cmp = this._comparator(data, node.data);

        dir = cmp > 0;

        // save found node
        if (cmp === 0) {
            found = node;
        }

        // push the red node down
        if (!is_red(node) && !is_red(node.get_child(dir))) {
            if (is_red(node.get_child(!dir))) {
                var sr = single_rotate(node, dir);
                p.set_child(last, sr);
                p = sr;
            } else if (!is_red(node.get_child(!dir))) {
                var sibling = p.get_child(!last);
                if (sibling !== null) {
                    if (!is_red(sibling.get_child(!last)) && !is_red(sibling.get_child(last))) {
                        // color flip
                        p.red = false;
                        sibling.red = true;
                        node.red = true;
                    } else {
                        var dir2 = gp.right === p;

                        if (is_red(sibling.get_child(last))) {
                            gp.set_child(dir2, double_rotate(p, last));
                        } else if (is_red(sibling.get_child(!last))) {
                            gp.set_child(dir2, single_rotate(p, last));
                        }

                        // ensure correct coloring
                        var gpc = gp.get_child(dir2);
                        gpc.red = true;
                        node.red = true;
                        gpc.left.red = false;
                        gpc.right.red = false;
                    }
                }
            }
        }
    }

    // replace and remove if found
    if (found !== null) {
        found.data = node.data;
        p.set_child(p.right === node, node.get_child(node.left === null));
        this.size--;
    }

    // update root and make it black
    this._root = head.right;
    if (this._root !== null) {
        this._root.red = false;
    }

    return found !== null;
};

function is_red(node) {
    return node !== null && node.red;
}

function single_rotate(root, dir) {
    var save = root.get_child(!dir);

    root.set_child(!dir, save.get_child(dir));
    save.set_child(dir, root);

    root.red = true;
    save.red = false;

    return save;
}

function double_rotate(root, dir) {
    root.set_child(!dir, single_rotate(root.get_child(!dir), !dir));
    return single_rotate(root, dir);
}

module.exports = RBTree;