|
- # cython: profile=False
- # cython: cdivision=True
-
- """
- Jim Paris <jim@jtan.com>
-
- 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.
-
- Code that assumes non-overlapping intervals is marked with the
- string 'non-overlapping'.
- """
-
- import sys
- cimport rbtree
-
- cdef class RBNode:
- """One node of the Red/Black tree, containing a key (start, end)
- and value (obj)"""
- def __init__(self, double start, double end, object obj = None):
- self.obj = obj
- self.start = start
- self.end = end
- self.red = False
- self.left = None
- self.right = None
-
- def __str__(self):
- if self.red:
- color = "R"
- else:
- color = "B"
- if self.start == sys.float_info.min:
- return "[node nil]"
- return ("[node ("
- + str(self.obj) + ") "
- + str(self.start) + " -> " + str(self.end) + " "
- + color + "]")
-
- cdef class RBTree:
- """Red/Black tree"""
-
- # Init
- def __init__(self):
- self.nil = RBNode(start = sys.float_info.min,
- end = sys.float_info.min)
- self.nil.left = self.nil
- self.nil.right = self.nil
- self.nil.parent = self.nil
-
- self.root = RBNode(start = sys.float_info.max,
- end = sys.float_info.max)
- self.root.left = self.nil
- 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.
- cpdef getroot(self):
- return self.root.left
-
- # Rotations and basic operations
- cdef void __rotate_left(self, RBNode x):
- """Rotate left:
- # x y
- # / \ --> / \
- # z y x w
- # / \ / \
- # v w z v
- """
- cdef RBNode y = x.right
- x.right = y.left
- if y.left is not self.nil:
- y.left.parent = x
- y.parent = x.parent
- if x is x.parent.left:
- x.parent.left = y
- else:
- x.parent.right = y
- y.left = x
- x.parent = y
-
- cdef void __rotate_right(self, RBNode y):
- """Rotate right:
- # y x
- # / \ --> / \
- # x w z y
- # / \ / \
- # z v v w
- """
- cdef RBNode x = y.left
- y.left = x.right
- if x.right is not self.nil:
- x.right.parent = y
- x.parent = y.parent
- if y is y.parent.left:
- y.parent.left = x
- else:
- y.parent.right = x
- x.right = y
- y.parent = x
-
- cdef RBNode __successor(self, RBNode x):
- """Returns the successor of RBNode x"""
- cdef RBNode y = x.right
- if y is not self.nil:
- while y.left is not self.nil:
- y = y.left
- else:
- y = x.parent
- while x is y.right:
- x = y
- y = y.parent
- if y is self.root:
- return self.nil
- return y
- cpdef RBNode successor(self, RBNode x):
- """Returns the successor of RBNode x, or None"""
- cdef RBNode y = self.__successor(x)
- return y if y is not self.nil else None
-
- cdef RBNode __predecessor(self, RBNode x):
- """Returns the predecessor of RBNode x"""
- cdef RBNode y = x.left
- if y is not self.nil:
- while y.right is not self.nil:
- y = y.right
- else:
- y = x.parent
- while x is y.left:
- if y is self.root:
- y = self.nil
- break
- x = y
- y = y.parent
- return y
- cpdef RBNode predecessor(self, RBNode x):
- """Returns the predecessor of RBNode x, or None"""
- cdef RBNode y = self.__predecessor(x)
- return y if y is not self.nil else None
-
- # Insertion
- cpdef insert(self, RBNode z):
- """Insert RBNode z into RBTree and rebalance as necessary"""
- z.left = self.nil
- z.right = self.nil
- cdef RBNode y = self.root
- cdef RBNode x = self.root.left
- while x is not self.nil:
- y = x
- if (x.start > z.start or (x.start == z.start and x.end > z.end)):
- x = x.left
- else:
- x = x.right
- z.parent = y
- if (y is self.root or
- (y.start > z.start or (y.start == z.start and y.end > z.end))):
- y.left = z
- else:
- y.right = z
- # relabel/rebalance
- self.__insert_fixup(z)
-
- cdef void __insert_fixup(self, RBNode x):
- """Rebalance/fix RBTree after a simple insertion of RBNode x"""
- x.red = True
- while x.parent.red:
- if x.parent is x.parent.parent.left:
- y = x.parent.parent.right
- if y.red:
- x.parent.red = False
- y.red = False
- x.parent.parent.red = True
- x = x.parent.parent
- else:
- if x is x.parent.right:
- x = x.parent
- self.__rotate_left(x)
- x.parent.red = False
- x.parent.parent.red = True
- self.__rotate_right(x.parent.parent)
- else: # same as above, left/right switched
- y = x.parent.parent.left
- if y.red:
- x.parent.red = False
- y.red = False
- x.parent.parent.red = True
- x = x.parent.parent
- else:
- if x is x.parent.left:
- x = x.parent
- self.__rotate_right(x)
- x.parent.red = False
- x.parent.parent.red = True
- self.__rotate_left(x.parent.parent)
- self.root.left.red = False
-
- # Deletion
- cpdef delete(self, RBNode z):
- if z.left is None or z.right is None:
- raise AttributeError("you can only delete a node object "
- + "from the tree; use find() to get one")
- cdef RBNode x, y
- if z.left is self.nil or z.right is self.nil:
- y = z
- else:
- y = self.__successor(z)
- if y.left is self.nil:
- x = y.right
- else:
- x = y.left
- x.parent = y.parent
- if x.parent is self.root:
- self.root.left = x
- else:
- if y is y.parent.left:
- y.parent.left = x
- else:
- y.parent.right = x
- if y is not z:
- # y is the node to splice out, x is its child
- y.left = z.left
- y.right = z.right
- y.parent = z.parent
- z.left.parent = y
- z.right.parent = y
- if z is z.parent.left:
- z.parent.left = y
- else:
- z.parent.right = y
- if not y.red:
- y.red = z.red
- self.__delete_fixup(x)
- else:
- y.red = z.red
- else:
- if not y.red:
- self.__delete_fixup(x)
-
- cdef void __delete_fixup(self, RBNode x):
- """Rebalance/fix RBTree after a deletion. RBNode x is the
- child of the spliced out node."""
- cdef RBNode rootLeft = self.root.left
- while not x.red and x is not rootLeft:
- if x is x.parent.left:
- w = x.parent.right
- if w.red:
- w.red = False
- x.parent.red = True
- self.__rotate_left(x.parent)
- w = x.parent.right
- if not w.right.red and not w.left.red:
- w.red = True
- x = x.parent
- else:
- if not w.right.red:
- w.left.red = False
- w.red = True
- self.__rotate_right(w)
- w = x.parent.right
- w.red = x.parent.red
- x.parent.red = False
- w.right.red = False
- self.__rotate_left(x.parent)
- x = rootLeft # exit loop
- else: # same as above, left/right switched
- w = x.parent.left
- if w.red:
- w.red = False
- x.parent.red = True
- self.__rotate_right(x.parent)
- w = x.parent.left
- if not w.left.red and not w.right.red:
- w.red = True
- x = x.parent
- else:
- if not w.left.red:
- w.right.red = False
- w.red = True
- self.__rotate_left(w)
- w = x.parent.left
- w.red = x.parent.red
- x.parent.red = False
- w.left.red = False
- self.__rotate_right(x.parent)
- x = rootLeft # exit loop
- x.red = False
-
- # Walking, searching
- def __iter__(self):
- return self.inorder()
-
- def inorder(self, RBNode x = None):
- """Generator that performs an inorder walk for the tree
- rooted at RBNode x"""
- if x is None:
- 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)
-
- cpdef RBNode find(self, double start, double end):
- """Return the node with exactly the given start and end."""
- cdef RBNode x = self.getroot()
- while x is not self.nil:
- if start < x.start:
- x = x.left
- elif start == x.start:
- if end == x.end:
- break # found it
- elif end < x.end:
- x = x.left
- else:
- x = x.right
- else:
- x = x.right
- return x if x is not self.nil else None
-
- cpdef RBNode find_left_end(self, double t):
- """Find the leftmode node with end >= t. With non-overlapping
- intervals, this is the first node that might overlap time t.
-
- 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."""
- cdef RBNode 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
-
- cpdef RBNode find_right_start(self, double t):
- """Find the rightmode node with start <= t. With non-overlapping
- intervals, this is the last node that might overlap time t."""
- cdef RBNode 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, double start, double end):
- """Generator that returns nodes that overlap the given
- (start,end) range. Assumes non-overlapping intervals."""
- # Start with the leftmode node that ends after start
- cdef RBNode n = self.find_left_end(start)
- while n is not None:
- if n.start >= end:
- # this node starts after the requested end; we're done
- break
- if start < n.end:
- # this node overlaps our requested area
- yield n
- n = self.successor(n)
|