Implement Order Statistic Tree (fixes #15)

README.md

@@ -68,6 +68,12 @@ Methods
 ### clear()
 > Removes all nodes from the tree.
 
+### get(index)
+> Get the item at the given index.
+
+### rank(item)
+> Get the index of the given item.
+
 ### find(item)
 > Returns node data if found, null otherwise.
 

lib/bintree.js

@@ -5,6 +5,7 @@ function Node(data) {
     this.data = data;
     this.left = null;
     this.right = null;
+    this.size = 1;
 }
 
 Node.prototype.get_child = function(dir) {
@@ -63,6 +64,7 @@ BinTree.prototype.insert = function(data) {
 
         // update helpers
         p = node;
+        node.size++;
         node = node.get_child(dir);
     }
 };
@@ -85,6 +87,7 @@ BinTree.prototype.remove = function(data) {
         node = node.get_child(dir);
         var cmp = this._comparator(data, node.data);
         dir = cmp > 0;
+        node.size--;
 
         if(cmp === 0) {
             found = node;

lib/rbtree.js

@@ -6,6 +6,7 @@ function Node(data) {
     this.left = null;
     this.right = null;
     this.red = true;
+    this.size = 1;
 }
 
 Node.prototype.get_child = function(dir) {
@@ -60,6 +61,14 @@ RBTree.prototype.insert = function(data) {
                 p.set_child(dir, node);
                 ret = true;
                 this.size++;
+
+                // update the sizes from root to the new node
+                var root = head.right;
+                while (root !== node) {
+                    root.size++;
+                    var c = this._comparator(data, root.data);
+                    root = root.get_child(c > 0);
+                }
             }
             else if(is_red(node.left) && is_red(node.right)) {
                 // color flip
@@ -130,6 +139,7 @@ RBTree.prototype.remove = function(data) {
         gp = p;
         p = node;
         node = node.get_child(dir);
+        node.size--;
 
         var cmp = this._comparator(data, node.data);
 
@@ -207,6 +217,10 @@ function single_rotate(root, dir) {
     root.red = true;
     save.red = false;
 
+    var child = save.get_child(!dir);
+    save.size = root.size;
+    root.size -= (child && child.size) + 1;
+
     return save;
 }
 

lib/treebase.js

@@ -139,6 +139,35 @@ TreeBase.prototype.reach = function(cb) {
     }
 };
 
+// get the data at the given index
+TreeBase.prototype.get = function(index) {
+    var select = function (node, i) {
+        if (!node) {
+            return null;
+        }
+        var r = (node.left && node.left.size) + 1;
+        return i === r ? node.data : i < r ? select(node.left, i) : select(node.right, i - r);
+    };
+    return select(this._root, index + 1);
+};
+
+// get the index of the given data
+TreeBase.prototype.rank = function(data) {
+    var result = 0;
+    var node = this._root;
+    while (node) {
+        var c = this._comparator(data, node.data);
+        if (c >= 0) {
+            result += (node.left && node.left.size) + (c > 0);
+            if (c === 0) {
+                break;
+            }
+        }
+        node = node.get_child(c > 0);
+    }
+    return result;
+};
+
 
 function Iterator(tree) {
     this._tree = tree;

test/test_api.js

@@ -48,6 +48,32 @@ function minmax(assert, tree_class) {
     assert.equal(tree.max(), _.max(inserts));
 }
 
+function get(assert, tree_class) {
+    var inserts = loader.get_inserts(loader.load(SAMPLE_FILE));
+    var tree = loader.build_tree(tree_class, inserts);
+    var remove = inserts.splice(0, inserts.length / 10);
+    remove.forEach(function(r) {
+        tree.remove(r);
+    });
+    inserts.sort(function(a,b) { return a - b; });
+    for(var i = 0; i < inserts.length; i++) {
+        assert.equal(tree.get(i), inserts[i]);
+    }
+}
+
+function rank(assert, tree_class) {
+    var inserts = loader.get_inserts(loader.load(SAMPLE_FILE));
+    var tree = loader.build_tree(tree_class, inserts);
+    var remove = inserts.splice(0, inserts.length / 10);
+    remove.forEach(function(r) {
+        tree.remove(r);
+    });
+    inserts.sort(function(a,b) { return a - b; });
+    for(var i = 0; i < inserts.length; i++) {
+        assert.equal(tree.rank(inserts[i]), i);
+    }
+}
+
 function forward_it(assert, tree_class) {
     var inserts = loader.get_inserts(loader.load(SAMPLE_FILE));
     var tree = loader.build_tree(tree_class, inserts);
@@ -294,6 +320,8 @@ var TESTS = {
     dup: dup,
     nonexist: nonexist,
     minmax: minmax,
+    get: get,
+    rank: rank,
     forward_it: forward_it,
     forward_it_break: forward_it_break,
     reverse_it: reverse_it,