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.
 
 
 

379 lines
12 KiB

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