379 lines
12 KiB
Cython
379 lines
12 KiB
Cython
# cython: profile=False
|
|
# cython: cdivision=True
|
|
# cython: language_level=2
|
|
|
|
"""
|
|
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)
|