Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with
or
.
Download ZIP
Browse files

Adds a getKthValue method to goog.structs.AvlTree that runs in O(log …

…n) time.

R=chrishenry
DELTA=163 (133 added, 21 deleted, 9 changed)


Revision created by MOE tool push_codebase.
MOE_MIGRATION=6265


git-svn-id: http://closure-library.googlecode.com/svn/trunk@2512 0b95b8e8-c90f-11de-9d4f-f947ee5921c8
  • Loading branch information...
commit 4955c46829489c38a26073bcd543ad30ee5d35fd 1 parent ab09db2
danchen@google.com authored
View
120 closure/goog/structs/avltree.js
@@ -125,15 +125,6 @@ goog.structs.AvlTree.prototype.maxNode_ = null;
/**
- * Keeps track of the number of nodes in the tree.
- *
- * @type {number}
- * @private
- */
-goog.structs.AvlTree.prototype.count_ = 0;
-
-
-/**
* Inserts a node into the tree with the specified value if the tree does
* not already contain a node with the specified value. If the value is
* inserted, the tree is balanced to enforce the AVL-Tree height property.
@@ -148,12 +139,11 @@ goog.structs.AvlTree.prototype.add = function(value) {
this.root_ = new goog.structs.AvlTree.Node(value);
this.minNode_ = this.root_;
this.maxNode_ = this.root_;
- this.count_ = 1;
return true;
}
- // Assume a node is not added and change status when one is
- var retStatus = false;
+ // This will be set to the new node if a new node is added.
+ var newNode = null;
// Depth traverse the tree and insert the value if we reach a null node
this.traverse_(function(node) {
@@ -161,36 +151,38 @@ goog.structs.AvlTree.prototype.add = function(value) {
if (this.comparator_(node.value, value) > 0) {
retNode = node.left;
if (node.left == null) {
- var newNode = new goog.structs.AvlTree.Node(value, node);
+ newNode = new goog.structs.AvlTree.Node(value, node);
node.left = newNode;
if (node == this.minNode_) {
this.minNode_ = newNode;
}
- retStatus = true; // Value was added to tree
- this.balance_(node); // Maintain the AVL-tree balance
}
} else if (this.comparator_(node.value, value) < 0) {
retNode = node.right;
if (node.right == null) {
- var newNode = new goog.structs.AvlTree.Node(value, node);
+ newNode = new goog.structs.AvlTree.Node(value, node);
node.right = newNode;
if (node == this.maxNode_) {
this.maxNode_ = newNode;
}
- retStatus = true; // Value was added to tree
- this.balance_(node); // Maintain the AVL-tree balance
}
}
return retNode; // If null, we'll stop traversing the tree
});
- // If a node was added, increment count
- if (retStatus) {
- this.count_ += 1;
+ // If a node was added, increment counts and balance tree.
+ if (newNode) {
+ this.traverse_(
+ function(node) {
+ node.count++;
+ return node.parent;
+ },
+ newNode.parent);
+ this.balance_(newNode.parent); // Maintain the AVL-tree balance
}
// Return true if a node was added, false otherwise
- return retStatus;
+ return !!newNode;
};
@@ -222,12 +214,6 @@ goog.structs.AvlTree.prototype.remove = function(value) {
return retNode; // If null, we'll stop traversing the tree
});
- // If a node was removed, decrement count.
- if (retValue) {
- // Had traverse_() cleared the tree, set to 0.
- this.count_ = this.root_ ? this.count_ - 1 : 0;
- }
-
// Return the value that was removed, null if the value was not in the tree
return retValue;
};
@@ -240,7 +226,6 @@ goog.structs.AvlTree.prototype.clear = function() {
this.root_ = null;
this.minNode_ = null;
this.maxNode_ = null;
- this.count_ = 0;
};
@@ -281,7 +266,21 @@ goog.structs.AvlTree.prototype.contains = function(value) {
* @override
*/
goog.structs.AvlTree.prototype.getCount = function() {
- return this.count_;
+ return this.root_ ? this.root_.count : 0;
+};
+
+
+/**
+ * Returns a k-th smallest value, based on the comparator, where 0 <= k <
+ * this.getCount().
+ * @param {number} k The number k.
+ * @return {*} The k-th smallest value.
+ */
+goog.structs.AvlTree.prototype.getKthValue = function(k) {
+ if (k < 0 || k >= this.getCount()) {
+ return null;
+ }
+ return this.getKthNode_(k).value;
};
@@ -544,6 +543,10 @@ goog.structs.AvlTree.prototype.leftRotate_ = function(node) {
if (node.right != null) node.right.parent = node;
temp.left = node;
node.parent = temp;
+
+ // Update counts.
+ temp.count = node.count;
+ node.count -= (temp.right ? temp.right.count : 0) + 1;
};
@@ -572,6 +575,10 @@ goog.structs.AvlTree.prototype.rightRotate_ = function(node) {
if (node.left != null) node.left.parent = node;
temp.right = node;
node.parent = temp;
+
+ // Update counts.
+ temp.count = node.count;
+ node.count -= (temp.left ? temp.left.count : 0) + 1;
};
@@ -590,6 +597,13 @@ goog.structs.AvlTree.prototype.removeNode_ = function(node) {
var r; // Node to replace the node being removed
if (node.left != null) {
r = this.getMaxNode_(node.left);
+
+ // Update counts.
+ this.traverse_(function(node) {
+ node.count--;
+ return node.parent;
+ }, r);
+
if (r != node.left) {
r.parent.right = r.left;
if (r.left) r.left.parent = r.parent;
@@ -601,8 +615,16 @@ goog.structs.AvlTree.prototype.removeNode_ = function(node) {
r.right = node.right;
if (r.right) r.right.parent = r;
if (node == this.maxNode_) this.maxNode_ = r;
+ r.count = node.count;
} else {
r = this.getMinNode_(node.right);
+
+ // Update counts.
+ this.traverse_(function(node) {
+ node.count--;
+ return node.parent;
+ }, r);
+
if (r != node.right) {
r.parent.left = r.right;
if (r.right) r.right.parent = r.parent;
@@ -614,6 +636,7 @@ goog.structs.AvlTree.prototype.removeNode_ = function(node) {
r.left = node.left;
if (r.left) r.left.parent = r;
if (node == this.minNode_) this.minNode_ = r;
+ r.count = node.count;
}
// Update the parent of the node being removed to point to its replace
@@ -628,6 +651,12 @@ goog.structs.AvlTree.prototype.removeNode_ = function(node) {
// Balance the tree
this.balance_(b ? b : r);
} else {
+ // Update counts.
+ this.traverse_(function(node) {
+ node.count--;
+ return node.parent;
+ }, node.parent);
+
// If the node is a leaf, remove it and balance starting from its parent
if (node.isLeftChild()) {
this.special = 1;
@@ -646,6 +675,30 @@ goog.structs.AvlTree.prototype.removeNode_ = function(node) {
/**
+ * Returns the node in the tree that has k nodes before it in an in-order
+ * traversal, optionally rooted at {@code opt_rootNode}.
+ *
+ * @param {number} k The number of nodes before the node to be returned in an
+ * in-order traversal, where 0 <= k < root.count.
+ * @param {goog.structs.AvlTree.Node=} opt_rootNode Optional root node.
+ * @return {goog.structs.AvlTree.Node} The node at the specified index.
+ * @private
+ */
+goog.structs.AvlTree.prototype.getKthNode_ = function(k, opt_rootNode) {
+ var root = opt_rootNode || this.root_;
+ var numNodesInLeftSubtree = root.left ? root.left.count : 0;
+
+ if (k < numNodesInLeftSubtree) {
+ return this.getKthNode_(k, root.left);
+ } else if (k == numNodesInLeftSubtree) {
+ return root;
+ } else {
+ return this.getKthNode_(k - numNodesInLeftSubtree - 1, root.right);
+ }
+};
+
+
+/**
* Returns the node with the smallest value in tree, optionally rooted at
* {@code opt_rootNode}.
*
@@ -725,6 +778,13 @@ goog.structs.AvlTree.Node = function(value, opt_parent) {
* @type {goog.structs.AvlTree.Node}
*/
this.parent = opt_parent ? opt_parent : null;
+
+ /**
+ * The number of nodes in the subtree rooted at this node.
+ *
+ * @type {number}
+ */
+ this.count = 1;
};
View
52 closure/goog/structs/avltree_test.html
@@ -244,8 +244,60 @@
tree.add(1);
tree.remove(1);
assertEquals(0, tree.getCount());
+
+ var values = ['bill', 'blake', 'elliot', 'jacob', 'john', 'myles', 'ted'];
+
+ // Insert strings into tree out of order
+ tree.add('frodo');
+ tree.add(values[4]);
+ tree.add(values[3]);
+ tree.add(values[0]);
+ tree.add(values[6]);
+ tree.add('samwise');
+ tree.add(values[5]);
+ tree.add(values[1]);
+ tree.add(values[2]);
+ tree.add('pippin');
+ assertEquals(10, tree.getCount());
+ assertEquals(
+ tree.root_.left.count + tree.root_.right.count + 1, tree.getCount());
+
+ // Remove strings from tree
+ assertEquals('samwise', tree.remove('samwise'));
+ assertEquals('pippin', tree.remove('pippin'));
+ assertEquals('frodo', tree.remove('frodo'));
+ assertEquals(null, tree.remove('merry'));
+ assertEquals(7, tree.getCount());
+
+ assertEquals(
+ tree.root_.left.count + tree.root_.right.count + 1, tree.getCount());
}
+ /**
+ * This test verifies that getKthOrder gets the correct value.
+ */
+ function testGetKthOrder() {
+ var tree = new goog.structs.AvlTree(function(a, b) {
+ return a - b;
+ });
+
+ var NUM_TO_INSERT = 2000;
+ var NUM_TO_REMOVE = 500;
+
+ // Insert ints into tree out of order
+ for (var i = 0; i < NUM_TO_INSERT; i += 1) {
+ tree.add(i);
+ }
+
+ // Remove values.
+ for (var i = 0; i < NUM_TO_REMOVE; i += 1) {
+ tree.remove(i);
+ }
+ for (var k = 0; k < tree.getCount(); ++k) {
+ assertEquals(NUM_TO_REMOVE + k, tree.getKthValue(k));
+ }
+ };
+
</script>
</body>
</html>
Please sign in to comment.
Something went wrong with that request. Please try again.