You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

378 lines
12 KiB

  1. # cython: profile=False
  2. # cython: cdivision=True
  3. """
  4. Jim Paris <jim@jtan.com>
  5. Red-black tree, where keys are stored as start/end timestamps.
  6. This is a basic interval tree that holds half-open intervals:
  7. [start, end)
  8. Intervals must not overlap. Fixing that would involve making this
  9. into an augmented interval tree as described in CLRS 14.3.
  10. Code that assumes non-overlapping intervals is marked with the
  11. string 'non-overlapping'.
  12. """
  13. import sys
  14. cimport rbtree
  15. cdef class RBNode:
  16. """One node of the Red/Black tree, containing a key (start, end)
  17. and value (obj)"""
  18. def __init__(self, double start, double end, object obj = None):
  19. self.obj = obj
  20. self.start = start
  21. self.end = end
  22. self.red = False
  23. self.left = None
  24. self.right = None
  25. def __str__(self):
  26. if self.red:
  27. color = "R"
  28. else:
  29. color = "B"
  30. if self.start == sys.float_info.min:
  31. return "[node nil]"
  32. return ("[node ("
  33. + str(self.obj) + ") "
  34. + str(self.start) + " -> " + str(self.end) + " "
  35. + color + "]")
  36. cdef class RBTree:
  37. """Red/Black tree"""
  38. # Init
  39. def __init__(self):
  40. self.nil = RBNode(start = sys.float_info.min,
  41. end = sys.float_info.min)
  42. self.nil.left = self.nil
  43. self.nil.right = self.nil
  44. self.nil.parent = self.nil
  45. self.root = RBNode(start = sys.float_info.max,
  46. end = sys.float_info.max)
  47. self.root.left = self.nil
  48. self.root.right = self.nil
  49. self.root.parent = self.nil
  50. # We have a dummy root node to simplify operations, so from an
  51. # external point of view, its left child is the real root.
  52. cpdef getroot(self):
  53. return self.root.left
  54. # Rotations and basic operations
  55. cdef void __rotate_left(self, RBNode x):
  56. """Rotate left:
  57. # x y
  58. # / \ --> / \
  59. # z y x w
  60. # / \ / \
  61. # v w z v
  62. """
  63. cdef RBNode y = x.right
  64. x.right = y.left
  65. if y.left is not self.nil:
  66. y.left.parent = x
  67. y.parent = x.parent
  68. if x is x.parent.left:
  69. x.parent.left = y
  70. else:
  71. x.parent.right = y
  72. y.left = x
  73. x.parent = y
  74. cdef void __rotate_right(self, RBNode y):
  75. """Rotate right:
  76. # y x
  77. # / \ --> / \
  78. # x w z y
  79. # / \ / \
  80. # z v v w
  81. """
  82. cdef RBNode x = y.left
  83. y.left = x.right
  84. if x.right is not self.nil:
  85. x.right.parent = y
  86. x.parent = y.parent
  87. if y is y.parent.left:
  88. y.parent.left = x
  89. else:
  90. y.parent.right = x
  91. x.right = y
  92. y.parent = x
  93. cdef RBNode __successor(self, RBNode x):
  94. """Returns the successor of RBNode x"""
  95. cdef RBNode y = x.right
  96. if y is not self.nil:
  97. while y.left is not self.nil:
  98. y = y.left
  99. else:
  100. y = x.parent
  101. while x is y.right:
  102. x = y
  103. y = y.parent
  104. if y is self.root:
  105. return self.nil
  106. return y
  107. cpdef RBNode successor(self, RBNode x):
  108. """Returns the successor of RBNode x, or None"""
  109. cdef RBNode y = self.__successor(x)
  110. return y if y is not self.nil else None
  111. cdef RBNode __predecessor(self, RBNode x):
  112. """Returns the predecessor of RBNode x"""
  113. cdef RBNode y = x.left
  114. if y is not self.nil:
  115. while y.right is not self.nil:
  116. y = y.right
  117. else:
  118. y = x.parent
  119. while x is y.left:
  120. if y is self.root:
  121. y = self.nil
  122. break
  123. x = y
  124. y = y.parent
  125. return y
  126. cpdef RBNode predecessor(self, RBNode x):
  127. """Returns the predecessor of RBNode x, or None"""
  128. cdef RBNode y = self.__predecessor(x)
  129. return y if y is not self.nil else None
  130. # Insertion
  131. cpdef insert(self, RBNode z):
  132. """Insert RBNode z into RBTree and rebalance as necessary"""
  133. z.left = self.nil
  134. z.right = self.nil
  135. cdef RBNode y = self.root
  136. cdef RBNode x = self.root.left
  137. while x is not self.nil:
  138. y = x
  139. if (x.start > z.start or (x.start == z.start and x.end > z.end)):
  140. x = x.left
  141. else:
  142. x = x.right
  143. z.parent = y
  144. if (y is self.root or
  145. (y.start > z.start or (y.start == z.start and y.end > z.end))):
  146. y.left = z
  147. else:
  148. y.right = z
  149. # relabel/rebalance
  150. self.__insert_fixup(z)
  151. cdef void __insert_fixup(self, RBNode x):
  152. """Rebalance/fix RBTree after a simple insertion of RBNode x"""
  153. x.red = True
  154. while x.parent.red:
  155. if x.parent is x.parent.parent.left:
  156. y = x.parent.parent.right
  157. if y.red:
  158. x.parent.red = False
  159. y.red = False
  160. x.parent.parent.red = True
  161. x = x.parent.parent
  162. else:
  163. if x is x.parent.right:
  164. x = x.parent
  165. self.__rotate_left(x)
  166. x.parent.red = False
  167. x.parent.parent.red = True
  168. self.__rotate_right(x.parent.parent)
  169. else: # same as above, left/right switched
  170. y = x.parent.parent.left
  171. if y.red:
  172. x.parent.red = False
  173. y.red = False
  174. x.parent.parent.red = True
  175. x = x.parent.parent
  176. else:
  177. if x is x.parent.left:
  178. x = x.parent
  179. self.__rotate_right(x)
  180. x.parent.red = False
  181. x.parent.parent.red = True
  182. self.__rotate_left(x.parent.parent)
  183. self.root.left.red = False
  184. # Deletion
  185. cpdef delete(self, RBNode z):
  186. if z.left is None or z.right is None:
  187. raise AttributeError("you can only delete a node object "
  188. + "from the tree; use find() to get one")
  189. cdef RBNode x, y
  190. if z.left is self.nil or z.right is self.nil:
  191. y = z
  192. else:
  193. y = self.__successor(z)
  194. if y.left is self.nil:
  195. x = y.right
  196. else:
  197. x = y.left
  198. x.parent = y.parent
  199. if x.parent is self.root:
  200. self.root.left = x
  201. else:
  202. if y is y.parent.left:
  203. y.parent.left = x
  204. else:
  205. y.parent.right = x
  206. if y is not z:
  207. # y is the node to splice out, x is its child
  208. y.left = z.left
  209. y.right = z.right
  210. y.parent = z.parent
  211. z.left.parent = y
  212. z.right.parent = y
  213. if z is z.parent.left:
  214. z.parent.left = y
  215. else:
  216. z.parent.right = y
  217. if not y.red:
  218. y.red = z.red
  219. self.__delete_fixup(x)
  220. else:
  221. y.red = z.red
  222. else:
  223. if not y.red:
  224. self.__delete_fixup(x)
  225. cdef void __delete_fixup(self, RBNode x):
  226. """Rebalance/fix RBTree after a deletion. RBNode x is the
  227. child of the spliced out node."""
  228. cdef RBNode rootLeft = self.root.left
  229. while not x.red and x is not rootLeft:
  230. if x is x.parent.left:
  231. w = x.parent.right
  232. if w.red:
  233. w.red = False
  234. x.parent.red = True
  235. self.__rotate_left(x.parent)
  236. w = x.parent.right
  237. if not w.right.red and not w.left.red:
  238. w.red = True
  239. x = x.parent
  240. else:
  241. if not w.right.red:
  242. w.left.red = False
  243. w.red = True
  244. self.__rotate_right(w)
  245. w = x.parent.right
  246. w.red = x.parent.red
  247. x.parent.red = False
  248. w.right.red = False
  249. self.__rotate_left(x.parent)
  250. x = rootLeft # exit loop
  251. else: # same as above, left/right switched
  252. w = x.parent.left
  253. if w.red:
  254. w.red = False
  255. x.parent.red = True
  256. self.__rotate_right(x.parent)
  257. w = x.parent.left
  258. if not w.left.red and not w.right.red:
  259. w.red = True
  260. x = x.parent
  261. else:
  262. if not w.left.red:
  263. w.right.red = False
  264. w.red = True
  265. self.__rotate_left(w)
  266. w = x.parent.left
  267. w.red = x.parent.red
  268. x.parent.red = False
  269. w.left.red = False
  270. self.__rotate_right(x.parent)
  271. x = rootLeft # exit loop
  272. x.red = False
  273. # Walking, searching
  274. def __iter__(self):
  275. return self.inorder()
  276. def inorder(self, RBNode x = None):
  277. """Generator that performs an inorder walk for the tree
  278. rooted at RBNode x"""
  279. if x is None:
  280. x = self.getroot()
  281. while x.left is not self.nil:
  282. x = x.left
  283. while x is not self.nil:
  284. yield x
  285. x = self.__successor(x)
  286. cpdef RBNode find(self, double start, double end):
  287. """Return the node with exactly the given start and end."""
  288. cdef RBNode x = self.getroot()
  289. while x is not self.nil:
  290. if start < x.start:
  291. x = x.left
  292. elif start == x.start:
  293. if end == x.end:
  294. break # found it
  295. elif end < x.end:
  296. x = x.left
  297. else:
  298. x = x.right
  299. else:
  300. x = x.right
  301. return x if x is not self.nil else None
  302. cpdef RBNode find_left_end(self, double t):
  303. """Find the leftmode node with end >= t. With non-overlapping
  304. intervals, this is the first node that might overlap time t.
  305. Note that this relies on non-overlapping intervals, since
  306. it assumes that we can use the endpoints to traverse the
  307. tree even though it was created using the start points."""
  308. cdef RBNode x = self.getroot()
  309. while x is not self.nil:
  310. if t < x.end:
  311. if x.left is self.nil:
  312. break
  313. x = x.left
  314. elif t == x.end:
  315. break
  316. else:
  317. if x.right is self.nil:
  318. x = self.__successor(x)
  319. break
  320. x = x.right
  321. return x if x is not self.nil else None
  322. cpdef RBNode find_right_start(self, double t):
  323. """Find the rightmode node with start <= t. With non-overlapping
  324. intervals, this is the last node that might overlap time t."""
  325. cdef RBNode x = self.getroot()
  326. while x is not self.nil:
  327. if t < x.start:
  328. if x.left is self.nil:
  329. x = self.__predecessor(x)
  330. break
  331. x = x.left
  332. elif t == x.start:
  333. break
  334. else:
  335. if x.right is self.nil:
  336. break
  337. x = x.right
  338. return x if x is not self.nil else None
  339. # Intersections
  340. def intersect(self, double start, double end):
  341. """Generator that returns nodes that overlap the given
  342. (start,end) range. Assumes non-overlapping intervals."""
  343. # Start with the leftmode node that ends after start
  344. cdef RBNode n = self.find_left_end(start)
  345. while n is not None:
  346. if n.start >= end:
  347. # this node starts after the requested end; we're done
  348. break
  349. if start < n.end:
  350. # this node overlaps our requested area
  351. yield n
  352. n = self.successor(n)