### Filling out rbtree tests, search routines

tags/before-insert-rework
Jim Paris 11 years ago
parent
commit
0b443f510b
3 changed files with 191 additions and 131 deletions
1. +74
-102
nilmdb/rbtree.py
2. +2
-2
setup.cfg
3. +115
-27
tests/test_rbtree.py

#### + 74 - 102 nilmdb/rbtree.pyView File

 @@ -3,6 +3,9 @@ This is a basic interval tree that holds half-open intervals: [start, end) Intervals must not overlap. Fixing that would involve making this into an augmented interval tree as described in CLRS 14.3. Code that assumes non-overlapping intervals is marked with the string 'non-overlapping'. """ import sys @@ -17,14 +20,17 @@ class RBNode(object): self.red = False self.left = None self.right = None self.nil = False def __str__(self): if self.red: color = "R" else: color = "B" return ("[node " + str(obj) + " " if self.start == sys.float_info.min: return "[node nil]" return ("[node (" + str(self.obj) + ") " + str(self.start) + " -> " + str(self.end) + " " + color + "]") @@ -38,7 +44,6 @@ class RBTree(object): self.nil.left = self.nil self.nil.right = self.nil self.nil.parent = self.nil self.nil.nil = True self.root = RBNode(start = sys.float_info.max, end = sys.float_info.max) @@ -46,6 +51,11 @@ class RBTree(object): self.root.right = self.nil self.root.parent = self.nil # We have a dummy root node to simplify operations, so from an # external point of view, its left child is the real root. def getroot(self): return self.root.left # Rotations and basic operations def __rotate_left(self, x): """Rotate left: @@ -88,7 +98,7 @@ class RBTree(object): y.parent = x def __successor(self, x): """Returns the successor of x""" """Returns the successor of RBNode x""" y = x.right if y is not self.nil: while y.left is not self.nil: @@ -101,6 +111,10 @@ class RBTree(object): if y is self.root: return self.nil return y def successor(self, x): """Returns the successor of RBNode x, or None""" y = self.__successor(x) return y if y is not self.nil else None def __predecessor(self, x): """Returns the predecessor of RBNode x""" @@ -117,6 +131,10 @@ class RBTree(object): x = y y = y.parent return y def predecessor(self, x): """Returns the predecessor of RBNode x, or None""" y = self.__predecessor(x) return y if y is not self.nil else None # Insertion def insert(self, z): @@ -264,129 +282,83 @@ class RBTree(object): x = rootLeft # exit loop x.red = False # Rendering def __render_dot_node(self, node, max_depth = 20): from printf import sprintf """Render a single node and its children into a dot graph fragment""" if max_depth == 0: return "" if node is self.nil: return "" def c(red): if red: return 'color="#ff0000", style=filled, fillcolor="#ffc0c0"' else: return 'color="#000000", style=filled, fillcolor="#c0c0c0"' s = sprintf("%d [label=\"%g\\n%g\", %s];\n", id(node), node.start, node.end, c(node.red)) if node.left is self.nil: s += sprintf("L%d [label=\"-\", %s];\n", id(node), c(False)) s += sprintf("%d -> L%d [label=L];\n", id(node), id(node)) else: s += sprintf("%d -> %d [label=L];\n", id(node), id(node.left)) if node.right is self.nil: s += sprintf("R%d [label=\"-\", %s];\n", id(node), c(False)) s += sprintf("%d -> R%d [label=R];\n", id(node), id(node)) else: s += sprintf("%d -> %d [label=R];\n", id(node), id(node.right)) s += self.__render_dot_node(node.left, max_depth-1) s += self.__render_dot_node(node.right, max_depth-1) return s def render_dot(self, title = "RBTree"): """Render the entire RBTree as a dot graph""" return ("digraph rbtree {\n" + self.__render_dot_node(self.root.left) + "}\n"); def render_dot_live(self, title = "RBTree"): """Render the entire RBTree as a dot graph, live GTK view""" import gtk import gtk.gdk sys.path.append("/usr/share/xdot") import xdot xdot.Pen.highlighted = lambda pen: pen s = ("digraph rbtree {\n" + self.__render_dot_node(self.root) + "}\n"); window = xdot.DotWindow() window.set_dotcode(s) window.set_title(title + " - any key to close") window.connect('destroy', gtk.main_quit) def quit(widget, event): if not event.is_modifier: window.destroy() gtk.main_quit() window.widget.connect('key-press-event', quit) gtk.main() # Walking, searching def __iter__(self): return self.inorder(self.root.left) return self.inorder() def inorder(self, x = None): """Generator that performs an inorder walk for the tree starting at RBNode x""" rooted at RBNode x""" if x is None: x = self.root.left x = self.getroot() while x.left is not self.nil: x = x.left while x is not self.nil: yield x x = self.__successor(x) def __find_all(self, start, end, x): """Find node with the specified (start,end) key. Also returns the largest node less than or equal to key, and the smallest node greater or equal to than key.""" if x is None: x = self.root.left largest = self.nil smallest = self.nil def find(self, start, end): """Return the node with exactly the given start and end.""" x = self.getroot() while x is not self.nil: if start < x.start: smallest = x x = x.left # start < x = x.left elif start == x.start: if end < x.end: smallest = x x = x.left # start =, end < elif end == x.end: # found it smallest = x largest = x break if end == x.end: break # found it elif end < x.end: x = x.left else: largest = x x = x.right # start =, end > x = x.right else: largest = x x = x.right # start > return (x, smallest, largest) x = x.right return x if x is not self.nil else None def find(self, start, end, x = None): """Find node with the key == (start,end), or None""" y = self.__find_all(start, end, x)[1] return y if y is not self.nil else None def find_left_end(self, t): """Find the leftmode node with end >= t. With non-overlapping intervals, this is the first node that might overlap time t. def find_right(self, start, end, x = None): """Find node with the smallest key >= (start,end), or None""" y = self.__find_all(start, end, x)[1] return y if y is not self.nil else None Note that this relies on non-overlapping intervals, since it assumes that we can use the endpoints to traverse the tree even though it was created using the start points.""" x = self.getroot() while x is not self.nil: if t < x.end: if x.left is self.nil: break x = x.left elif t == x.end: break else: if x.right is self.nil: x = self.__successor(x) break x = x.right return x if x is not self.nil else None def find_left(self, start, end, x = None): """Find node with the largest key <= (start,end), or None""" y = self.__find_all(start, end, x)[2] return y if y is not self.nil else None def find_right_start(self, t): """Find the rightmode node with start <= t. With non-overlapping intervals, this is the last node that might overlap time t.""" x = self.getroot() while x is not self.nil: if t < x.start: if x.left is self.nil: x = self.__predecessor(x) break x = x.left elif t == x.start: break else: if x.right is self.nil: break x = x.right return x if x is not self.nil else None # Intersections def intersect(self, start, end): """Generator that returns nodes that overlap the given (start,end) range, for the tree rooted at RBNode x. (start,end) range. Assumes non-overlapping intervals.""" NOTE: this assumes non-overlapping intervals.""" # Start with the leftmost node before the starting point n = self.find_left(start, start) # If we didn't find one, look for the leftmode node before the

#### + 2 - 2 setup.cfgView File

 @@ -12,8 +12,8 @@ stop= verbosity=2 #tests=tests/test_cmdline.py #tests=tests/test_layout.py #tests=tests/test_rbtree.py tests=tests/test_interval.py tests=tests/test_rbtree.py #tests=tests/test_interval.py #tests=tests/test_client.py #tests=tests/test_timestamper.py #tests=tests/test_serializer.py

#### + 115 - 27 tests/test_rbtree.pyView File

 @@ -11,17 +11,33 @@ from nilmdb.rbtree import RBTree, RBNode from test_helpers import * import unittest render = False # set to False to skip live renders do_live_renders = False def render(tree, description = "", live = True): import renderdot if not do_live_renders: # If not doing a live render, still render it to a string. live = False r = renderdot.Renderer(lambda node: node.left, lambda node: node.right, lambda node: node.red, lambda node: node.start, lambda node: node.end, tree.nil) if live: r.render_dot_live(tree.getroot(), description) else: return r.render_dot(tree.getroot()) class TestRBTree: def test_rbtree(self): rb = RBTree() rb.insert(RBNode(None, 10000, 10001)) rb.insert(RBNode(None, 10004, 10007)) rb.insert(RBNode(None, 10001, 10002)) s = rb.render_dot() rb.insert(RBNode(10000, 10001)) rb.insert(RBNode(10004, 10007)) rb.insert(RBNode(10001, 10002)) # There was a typo that gave the RBTree a loop in this case. # Verify that the dot isn't too big. s = render(rb, live = False) assert(len(s.splitlines()) < 30) def test_rbtree_big(self): @@ -32,44 +48,116 @@ class TestRBTree: rb = RBTree() j = 500 for i in xrange(j): rb.insert(RBNode(None, i, i+1)) # show the graph if render: rb.render_dot_live("in-order insert") rb.insert(RBNode(i, i+1)) render(rb, "in-order insert") # remove about half of them for i in random.sample(xrange(j),j): if random.randint(0,1): rb.delete(rb.find(i, i+1)) # show the graph if render: rb.render_dot_live("in-order insert, random delete") render(rb, "in-order insert, random delete") # make a set of 500 intervals, inserted at random rb = RBTree() j = 500 for i in random.sample(xrange(j),j): rb.insert(RBNode(None, i, i+1)) # show the graph if render: rb.render_dot_live("random insert") rb.insert(RBNode(i, i+1)) render(rb, "random insert") # remove about half of them for i in random.sample(xrange(j),j): if random.randint(0,1): rb.delete(rb.find(i, i+1)) # show the graph if render: rb.render_dot_live("random insert, random delete") render(rb, "random insert, random delete") # in-order insert of 250 more for i in xrange(250): rb.insert(RBNode(None, i+500, i+501)) rb.insert(RBNode(i+500, i+501)) render(rb, "random insert, random delete, in-order insert") # show the graph if render: rb.render_dot_live("random insert, random delete, in-order insert") def test_rbtree_basics(self): rb = RBTree() vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4] for n in vals: rb.insert(RBNode(n, n)) # stringify s = "" for node in rb: s += str(node) assert "[node (None) 1 -> 1 B]" in s assert str(rb.nil) == "[node nil]" # inorder traversal, successor and predecessor last = 0 for node in rb: assert(node.start > last) last = node.start successor = rb.successor(node) if successor: assert(rb.predecessor(successor) is node) predecessor = rb.predecessor(node) if predecessor: assert(rb.successor(predecessor) is node) # Delete node not in the tree with assert_raises(AttributeError): rb.delete(RBNode(1,2)) # Delete all nodes! for node in rb: rb.delete(node) # Build it up again, make sure it matches for n in vals: rb.insert(RBNode(n, n)) s2 = "" for node in rb: s2 += str(node) assert(s == s2) def test_rbtree_find(self): # Get a little bit of coverage for some overlapping cases, # even though the class doesn't fully support it. rb = RBTree() nodes = [ RBNode(1, 5), RBNode(1, 10), RBNode(1, 15) ] for n in nodes: rb.insert(n) assert(rb.find(1, 5) is nodes[0]) assert(rb.find(1, 10) is nodes[1]) assert(rb.find(1, 15) is nodes[2]) def test_rbtree_find_leftright(self): # Now let's get some ranges in there rb = RBTree() vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4] for n in vals: rb.insert(RBNode(n*10, n*10+5)) # Check find_end_left, find_right_start for i in range(160): left = rb.find_left_end(i) right = rb.find_right_start(i) if left: # endpoint should be more than i assert(left.end >= i) # all earlier nodes should have a lower endpoint for node in rb: if node is left: break assert(node.end < i) if right: # startpoint should be less than i assert(right.start <= i) # all later nodes should have a higher startpoint for node in reversed(list(rb)): if node is right: break assert(node.start > i) def test_rbtree_intersect(self): # Fill with some ranges again rb = RBTree() vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4] for n in vals: rb.insert(RBNode(n*10, n*10+5))