Some documentation and other cleanups in rbtree.py

nilmdb/rbtree.py

@@ -1,20 +1,17 @@
-"""Red-black tree, where keys are stored as start/end timestamps."""
+"""Red-black tree, where keys are stored as start/end timestamps.
+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.
+"""
 
 import sys
 
 class RBNode(object):
-    """One node of the Red/Black tree.  obj points to any object,
-    'start' and 'end' are timestamps that represent the key."""
-    def __init__(self, obj = None, start = None, end = None):
-        """If given an object but no start/end times, get the
-        start/end times from the object.
-
-        If given start/end times, obj can be anything, including None."""
+    """One node of the Red/Black tree, containing a key (start, end)
+    and value (obj)"""
+    def __init__(self, start, end, obj = None):
         self.obj = obj
-        if start is None:
-            start = obj.start
-        if end is None:
-            end = obj.end
         self.start = start
         self.end = end
         self.red = False
@@ -27,6 +24,7 @@ class RBNode(object):
         else:
             color = "B"
         return ("[node "
+                + str(obj) + " "
                 + str(self.start) + " -> " + str(self.end) + " "
                 + color + "]")
 
@@ -36,7 +34,7 @@ class RBTree(object):
     # Init
     def __init__(self):
         self.nil = RBNode(start = sys.float_info.min,
-                           end = sys.float_info.min)
+                          end = sys.float_info.min)
         self.nil.left = self.nil
         self.nil.right = self.nil
         self.nil.parent = self.nil
@@ -50,6 +48,13 @@ class RBTree(object):
 
     # Rotations and basic operations
     def __rotate_left(self, x):
+        """Rotate left:
+        #   x           y
+        #  / \   -->   / \
+        # z   y       x   w
+        #    / \     / \
+        #   v   w   z   v
+        """
         y = x.right
         x.right = y.left
         if y.left is not self.nil:
@@ -63,6 +68,13 @@ class RBTree(object):
         x.parent = y
 
     def __rotate_right(self, y):
+        """Rotate right:
+        #     y           x
+        #    / \   -->   / \
+        #   x   w       z   y
+        #  / \             / \
+        # z   v           v   w
+        """
         x = y.left
         y.left = x.right
         if x.right is not self.nil:
@@ -76,7 +88,7 @@ class RBTree(object):
         y.parent = x
 
     def __successor(self, x):
-        """Returns the successor of RBNode x"""
+        """Returns the successor of x"""
         y = x.right
         if y is not self.nil:
             while y.left is not self.nil:
@@ -90,7 +102,7 @@ class RBTree(object):
                 return self.nil
         return y
 
-    def _predecessor(self, x):
+    def __predecessor(self, x):
         """Returns the predecessor of RBNode x"""
         y = x.left
         if y is not self.nil:
@@ -383,7 +395,7 @@ class RBTree(object):
             n = self.find_left(end, end)
         # If we still didn't find it, there are no intervals that intersect.
         if n is None:
-            return none
+            return
 
         # Now yield this node and all successors until their endpoints