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.

rbtree.pyx 12 KiB

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