|
|
@@ -1,4 +1,10 @@ |
|
|
|
"""Red-black tree, where keys are stored as start/end timestamps. |
|
|
|
# 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 |
|
|
@@ -10,17 +16,21 @@ string 'non-overlapping'. |
|
|
|
|
|
|
|
import sys |
|
|
|
|
|
|
|
class RBNode(object): |
|
|
|
cdef class RBNode: |
|
|
|
"""One node of the Red/Black tree, containing a key (start, end) |
|
|
|
and value (obj)""" |
|
|
|
def __init__(self, start, end, obj = None): |
|
|
|
cdef public object obj |
|
|
|
cdef public double start, end |
|
|
|
cdef public int red |
|
|
|
cdef public RBNode left, right, parent |
|
|
|
|
|
|
|
def __cinit__(RBNode 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 |
|
|
|
self.nil = False |
|
|
|
|
|
|
|
def __str__(self): |
|
|
|
if self.red: |
|
|
@@ -34,11 +44,13 @@ class RBNode(object): |
|
|
|
+ str(self.start) + " -> " + str(self.end) + " " |
|
|
|
+ color + "]") |
|
|
|
|
|
|
|
class RBTree(object): |
|
|
|
cdef class RBTree: |
|
|
|
"""Red/Black tree""" |
|
|
|
|
|
|
|
cdef public RBNode nil, root |
|
|
|
|
|
|
|
# Init |
|
|
|
def __init__(self): |
|
|
|
def __cinit__(RBTree self): |
|
|
|
self.nil = RBNode(start = sys.float_info.min, |
|
|
|
end = sys.float_info.min) |
|
|
|
self.nil.left = self.nil |
|
|
@@ -53,11 +65,11 @@ class RBTree(object): |
|
|
|
|
|
|
|
# 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): |
|
|
|
cpdef getroot(RBTree self): |
|
|
|
return self.root.left |
|
|
|
|
|
|
|
# Rotations and basic operations |
|
|
|
def __rotate_left(self, x): |
|
|
|
cdef void __rotate_left(RBTree self, RBNode x): |
|
|
|
"""Rotate left: |
|
|
|
# x y |
|
|
|
# / \ --> / \ |
|
|
@@ -65,7 +77,7 @@ class RBTree(object): |
|
|
|
# / \ / \ |
|
|
|
# v w z v |
|
|
|
""" |
|
|
|
y = x.right |
|
|
|
cdef RBNode y = x.right |
|
|
|
x.right = y.left |
|
|
|
if y.left is not self.nil: |
|
|
|
y.left.parent = x |
|
|
@@ -77,7 +89,7 @@ class RBTree(object): |
|
|
|
y.left = x |
|
|
|
x.parent = y |
|
|
|
|
|
|
|
def __rotate_right(self, y): |
|
|
|
cdef void __rotate_right(RBTree self, RBNode y): |
|
|
|
"""Rotate right: |
|
|
|
# y x |
|
|
|
# / \ --> / \ |
|
|
@@ -85,7 +97,7 @@ class RBTree(object): |
|
|
|
# / \ / \ |
|
|
|
# z v v w |
|
|
|
""" |
|
|
|
x = y.left |
|
|
|
cdef RBNode x = y.left |
|
|
|
y.left = x.right |
|
|
|
if x.right is not self.nil: |
|
|
|
x.right.parent = y |
|
|
@@ -97,9 +109,9 @@ class RBTree(object): |
|
|
|
x.right = y |
|
|
|
y.parent = x |
|
|
|
|
|
|
|
def __successor(self, x): |
|
|
|
cdef RBNode __successor(RBTree self, RBNode x): |
|
|
|
"""Returns the successor of RBNode x""" |
|
|
|
y = x.right |
|
|
|
cdef RBNode y = x.right |
|
|
|
if y is not self.nil: |
|
|
|
while y.left is not self.nil: |
|
|
|
y = y.left |
|
|
@@ -111,14 +123,14 @@ class RBTree(object): |
|
|
|
if y is self.root: |
|
|
|
return self.nil |
|
|
|
return y |
|
|
|
def successor(self, x): |
|
|
|
cpdef RBNode successor(RBTree self, RBNode x): |
|
|
|
"""Returns the successor of RBNode x, or None""" |
|
|
|
y = self.__successor(x) |
|
|
|
cdef RBNode y = self.__successor(x) |
|
|
|
return y if y is not self.nil else None |
|
|
|
|
|
|
|
def __predecessor(self, x): |
|
|
|
cdef RBNode __predecessor(RBTree self, RBNode x): |
|
|
|
"""Returns the predecessor of RBNode x""" |
|
|
|
y = x.left |
|
|
|
cdef RBNode y = x.left |
|
|
|
if y is not self.nil: |
|
|
|
while y.right is not self.nil: |
|
|
|
y = y.right |
|
|
@@ -131,18 +143,18 @@ class RBTree(object): |
|
|
|
x = y |
|
|
|
y = y.parent |
|
|
|
return y |
|
|
|
def predecessor(self, x): |
|
|
|
cpdef RBNode predecessor(RBTree self, RBNode x): |
|
|
|
"""Returns the predecessor of RBNode x, or None""" |
|
|
|
y = self.__predecessor(x) |
|
|
|
cdef RBNode y = self.__predecessor(x) |
|
|
|
return y if y is not self.nil else None |
|
|
|
|
|
|
|
# Insertion |
|
|
|
def insert(self, z): |
|
|
|
cpdef insert(RBTree self, RBNode z): |
|
|
|
"""Insert RBNode z into RBTree and rebalance as necessary""" |
|
|
|
z.left = self.nil |
|
|
|
z.right = self.nil |
|
|
|
y = self.root |
|
|
|
x = self.root.left |
|
|
|
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)): |
|
|
@@ -158,7 +170,7 @@ class RBTree(object): |
|
|
|
# relabel/rebalance |
|
|
|
self.__insert_fixup(z) |
|
|
|
|
|
|
|
def __insert_fixup(self, x): |
|
|
|
cdef void __insert_fixup(RBTree self, RBNode x): |
|
|
|
"""Rebalance/fix RBTree after a simple insertion of RBNode x""" |
|
|
|
x.red = True |
|
|
|
while x.parent.red: |
|
|
@@ -193,10 +205,11 @@ class RBTree(object): |
|
|
|
self.root.left.red = False |
|
|
|
|
|
|
|
# Deletion |
|
|
|
def delete(self, z): |
|
|
|
cpdef delete(RBTree 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: |
|
|
@@ -233,10 +246,10 @@ class RBTree(object): |
|
|
|
if not y.red: |
|
|
|
self.__delete_fixup(x) |
|
|
|
|
|
|
|
def __delete_fixup(self, x): |
|
|
|
cdef void __delete_fixup(RBTree self, RBNode x): |
|
|
|
"""Rebalance/fix RBTree after a deletion. RBNode x is the |
|
|
|
child of the spliced out node.""" |
|
|
|
rootLeft = self.root.left |
|
|
|
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 |
|
|
@@ -283,10 +296,10 @@ class RBTree(object): |
|
|
|
x.red = False |
|
|
|
|
|
|
|
# Walking, searching |
|
|
|
def __iter__(self): |
|
|
|
def __iter__(RBTree self): |
|
|
|
return self.inorder() |
|
|
|
|
|
|
|
def inorder(self, x = None): |
|
|
|
def inorder(RBTree self, RBNode x = None): |
|
|
|
"""Generator that performs an inorder walk for the tree |
|
|
|
rooted at RBNode x""" |
|
|
|
if x is None: |
|
|
@@ -297,9 +310,9 @@ class RBTree(object): |
|
|
|
yield x |
|
|
|
x = self.__successor(x) |
|
|
|
|
|
|
|
def find(self, start, end): |
|
|
|
cpdef RBNode find(RBTree self, double start, double end): |
|
|
|
"""Return the node with exactly the given start and end.""" |
|
|
|
x = self.getroot() |
|
|
|
cdef RBNode x = self.getroot() |
|
|
|
while x is not self.nil: |
|
|
|
if start < x.start: |
|
|
|
x = x.left |
|
|
@@ -314,14 +327,14 @@ class RBTree(object): |
|
|
|
x = x.right |
|
|
|
return x if x is not self.nil else None |
|
|
|
|
|
|
|
def find_left_end(self, t): |
|
|
|
cpdef RBNode find_left_end(RBTree 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.""" |
|
|
|
x = self.getroot() |
|
|
|
cdef RBNode x = self.getroot() |
|
|
|
while x is not self.nil: |
|
|
|
if t < x.end: |
|
|
|
if x.left is self.nil: |
|
|
@@ -336,10 +349,10 @@ class RBTree(object): |
|
|
|
x = x.right |
|
|
|
return x if x is not self.nil else None |
|
|
|
|
|
|
|
def find_right_start(self, t): |
|
|
|
cpdef RBNode find_right_start(RBTree self, double 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() |
|
|
|
cdef RBNode x = self.getroot() |
|
|
|
while x is not self.nil: |
|
|
|
if t < x.start: |
|
|
|
if x.left is self.nil: |
|
|
@@ -355,11 +368,11 @@ class RBTree(object): |
|
|
|
return x if x is not self.nil else None |
|
|
|
|
|
|
|
# Intersections |
|
|
|
def intersect(self, start, end): |
|
|
|
def intersect(RBTree 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 |
|
|
|
n = self.find_left_end(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 |