enumerating from an index/value in ruby

Question

So ive got this big-ass collection right? Its tree-backed(RBTree), so look-ups are fast, and the values are sorted.

Say ive got value X. I can look up X in my tree in logn time, cool. But I want the values to the right of X in the tree as well, until one of them dosent satisfy a test. Ie, get all elements that are >= X and < Y

I could get the index of X, i, then get the value at i+1, i+2.... until val(i+z) is > Y. Except that costs z * logn. If you enumerate over the tree, the enumerator inside the tree doesnt cost logn to proceed to the next element - it just follows a pointer in the tree.

So, is there a way to start enumerating from a specific index? Such that I dont have to skip over i elements before i can start enumerating over the range I want.

Tell me if you think im crazy.

Answer 1

Well, if you put the elements of your collection in there yourself, and you don't mind sorting them before insertion, you could wrap them with a linked list type. Just make sure the key for the linked list item wrapper for each element uses the element's key as its key for the tree sort. Then a lookup would get you a location in the linked list, and you could just walk it from there.

However, if you can't do it that way, your only recourse is to modify RBTree, which would require a little work in C, since it is a native extension to ruby. Most of the pieces that you need are there already dict_lookup() to give you the node in the tree you need, and rbtree_for_each() to show you how to write an iterator, given a start node.

You'd have to add the following code to rbtree.c in the RBTree gem:

*** rbtree.c.orig 2009-03-27 14:14:55.000000000 -0400
--- rbtree.c  2009-03-27 14:20:21.000000000 -0400
***************
*** 528,533 ****
--- 528,574 ----
      return EACH_NEXT;
  }

+ static VALUE
+ rbtree_each_starting_with_body(rbtree_each_arg_t* arg)
+ {
+     VALUE self = arg->self;
+     dict_t* dict = DICT(self);
+     dnode_t* node;
+     
+     ITER_LEV(self)++;
+     for (node = (dnode_t*) arg->arg;
+          node != NULL;
+          node = dict_next(dict, node)) {
+         
+         if (arg->func(node, NULL) == EACH_STOP)
+             break;
+     }
+     return self;
+ }
+ 
+ /*
+  * call-seq:
+  *   rbtree.each_starting_with(key) {|key, value| block} => rbtree
+  *
+  * Calls block once for each key in order, starting with the given key,
+  * passing the key and value as a two-element array parameters.
+  */
+ VALUE
+ rbtree_each_starting_with(VALUE self, VALUE key)
+ {
+     dnode_t* node = dict_lookup(DICT(self), TO_KEY(key));
+     rbtree_each_arg_t each_arg;
+     if (node == NULL) { return self; };
+     
+     RETURN_ENUMERATOR(self, 0, NULL);
+ 
+     each_arg.self = self;
+     each_arg.func = each_i;
+     each_arg.arg = node;
+     return rb_ensure(rbtree_each_starting_with_body, (VALUE)&each_arg,
+                      rbtree_each_ensure, self);
+ }
+ 
  /*
   * call-seq:
   *   rbtree.each {|key, value| block} => rbtree
***************
*** 1616,1621 ****
--- 1657,1663 ----
      rb_define_method(MultiRBTree, "length", rbtree_size, 0);

      rb_define_method(MultiRBTree, "each", rbtree_each, 0);
+     rb_define_method(MultiRBTree, "each_starting_with", rbtree_each_starting_with, 1);
      rb_define_method(MultiRBTree, "each_value", rbtree_each_value, 0);
      rb_define_method(MultiRBTree, "each_key", rbtree_each_key, 0);
      rb_define_method(MultiRBTree, "each_pair", rbtree_each_pair, 0);

Then if you run make in the source directory of the installed rbtree gem, it should remake the extension, and you can use it as normal:

% irb
irb> require 'rubygems'
=> true
irb> require 'rbtree'
=> true
irb> x = RBTree[ 'a', 1, 'b', 2, 'c', 3, 'd', 4 ]
=> #<RBTree: {"a"=>1, "b"=>2, "c"=>3, "d"=>4}, default=nil, cmp_proc=nil>
irb> x.each { |k,v| p [k, v] }
["a", 1]
["b", 2]
["c", 3]
["d", 4]
=> #<RBTree: {"a"=>1, "b"=>2, "c"=>3, "d"=>4}, default=nil, cmp_proc=nil>
irb> x.each_starting_with('b') { |k,v| p [k, v] }
["b", 2]
["c", 3]
["d", 4]
=> #<RBTree: {"a"=>1, "b"=>2, "c"=>3, "d"=>4}, default=nil, cmp_proc=nil>

Just remember that you've made this change, and distribute the modified gem with your changes. Or, hey, submit them to the gem creator on Rubyforge, so that everyone can take advantage of them.