interval_tree.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951
  1. // Copyright 2016 The etcd Authors
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package adt
  15. import (
  16. "bytes"
  17. "fmt"
  18. "math"
  19. "strings"
  20. )
  21. // Comparable is an interface for trichotomic comparisons.
  22. type Comparable interface {
  23. // Compare gives the result of a 3-way comparison
  24. // a.Compare(b) = 1 => a > b
  25. // a.Compare(b) = 0 => a == b
  26. // a.Compare(b) = -1 => a < b
  27. Compare(c Comparable) int
  28. }
  29. type rbcolor int
  30. const (
  31. black rbcolor = iota
  32. red
  33. )
  34. func (c rbcolor) String() string {
  35. switch c {
  36. case black:
  37. return "black"
  38. case red:
  39. return "black"
  40. default:
  41. panic(fmt.Errorf("unknown color %d", c))
  42. }
  43. }
  44. // Interval implements a Comparable interval [begin, end)
  45. // TODO: support different sorts of intervals: (a,b), [a,b], (a, b]
  46. type Interval struct {
  47. Begin Comparable
  48. End Comparable
  49. }
  50. // Compare on an interval gives == if the interval overlaps.
  51. func (ivl *Interval) Compare(c Comparable) int {
  52. ivl2 := c.(*Interval)
  53. ivbCmpBegin := ivl.Begin.Compare(ivl2.Begin)
  54. ivbCmpEnd := ivl.Begin.Compare(ivl2.End)
  55. iveCmpBegin := ivl.End.Compare(ivl2.Begin)
  56. // ivl is left of ivl2
  57. if ivbCmpBegin < 0 && iveCmpBegin <= 0 {
  58. return -1
  59. }
  60. // iv is right of iv2
  61. if ivbCmpEnd >= 0 {
  62. return 1
  63. }
  64. return 0
  65. }
  66. type intervalNode struct {
  67. // iv is the interval-value pair entry.
  68. iv IntervalValue
  69. // max endpoint of all descendent nodes.
  70. max Comparable
  71. // left and right are sorted by low endpoint of key interval
  72. left, right *intervalNode
  73. // parent is the direct ancestor of the node
  74. parent *intervalNode
  75. c rbcolor
  76. }
  77. func (x *intervalNode) color(sentinel *intervalNode) rbcolor {
  78. if x == sentinel {
  79. return black
  80. }
  81. return x.c
  82. }
  83. func (x *intervalNode) height(sentinel *intervalNode) int {
  84. if x == sentinel {
  85. return 0
  86. }
  87. ld := x.left.height(sentinel)
  88. rd := x.right.height(sentinel)
  89. if ld < rd {
  90. return rd + 1
  91. }
  92. return ld + 1
  93. }
  94. func (x *intervalNode) min(sentinel *intervalNode) *intervalNode {
  95. for x.left != sentinel {
  96. x = x.left
  97. }
  98. return x
  99. }
  100. // successor is the next in-order node in the tree
  101. func (x *intervalNode) successor(sentinel *intervalNode) *intervalNode {
  102. if x.right != sentinel {
  103. return x.right.min(sentinel)
  104. }
  105. y := x.parent
  106. for y != sentinel && x == y.right {
  107. x = y
  108. y = y.parent
  109. }
  110. return y
  111. }
  112. // updateMax updates the maximum values for a node and its ancestors
  113. func (x *intervalNode) updateMax(sentinel *intervalNode) {
  114. for x != sentinel {
  115. oldmax := x.max
  116. max := x.iv.Ivl.End
  117. if x.left != sentinel && x.left.max.Compare(max) > 0 {
  118. max = x.left.max
  119. }
  120. if x.right != sentinel && x.right.max.Compare(max) > 0 {
  121. max = x.right.max
  122. }
  123. if oldmax.Compare(max) == 0 {
  124. break
  125. }
  126. x.max = max
  127. x = x.parent
  128. }
  129. }
  130. type nodeVisitor func(n *intervalNode) bool
  131. // visit will call a node visitor on each node that overlaps the given interval
  132. func (x *intervalNode) visit(iv *Interval, sentinel *intervalNode, nv nodeVisitor) bool {
  133. if x == sentinel {
  134. return true
  135. }
  136. v := iv.Compare(&x.iv.Ivl)
  137. switch {
  138. case v < 0:
  139. if !x.left.visit(iv, sentinel, nv) {
  140. return false
  141. }
  142. case v > 0:
  143. maxiv := Interval{x.iv.Ivl.Begin, x.max}
  144. if maxiv.Compare(iv) == 0 {
  145. if !x.left.visit(iv, sentinel, nv) || !x.right.visit(iv, sentinel, nv) {
  146. return false
  147. }
  148. }
  149. default:
  150. if !x.left.visit(iv, sentinel, nv) || !nv(x) || !x.right.visit(iv, sentinel, nv) {
  151. return false
  152. }
  153. }
  154. return true
  155. }
  156. // IntervalValue represents a range tree node that contains a range and a value.
  157. type IntervalValue struct {
  158. Ivl Interval
  159. Val interface{}
  160. }
  161. // IntervalTree represents a (mostly) textbook implementation of the
  162. // "Introduction to Algorithms" (Cormen et al, 3rd ed.) chapter 13 red-black tree
  163. // and chapter 14.3 interval tree with search supporting "stabbing queries".
  164. type IntervalTree interface {
  165. // Insert adds a node with the given interval into the tree.
  166. Insert(ivl Interval, val interface{})
  167. // Delete removes the node with the given interval from the tree, returning
  168. // true if a node is in fact removed.
  169. Delete(ivl Interval) bool
  170. // Len gives the number of elements in the tree.
  171. Len() int
  172. // Height is the number of levels in the tree; one node has height 1.
  173. Height() int
  174. // MaxHeight is the expected maximum tree height given the number of nodes.
  175. MaxHeight() int
  176. // Visit calls a visitor function on every tree node intersecting the given interval.
  177. // It will visit each interval [x, y) in ascending order sorted on x.
  178. Visit(ivl Interval, ivv IntervalVisitor)
  179. // Find gets the IntervalValue for the node matching the given interval
  180. Find(ivl Interval) *IntervalValue
  181. // Intersects returns true if there is some tree node intersecting the given interval.
  182. Intersects(iv Interval) bool
  183. // Contains returns true if the interval tree's keys cover the entire given interval.
  184. Contains(ivl Interval) bool
  185. // Stab returns a slice with all elements in the tree intersecting the interval.
  186. Stab(iv Interval) []*IntervalValue
  187. // Union merges a given interval tree into the receiver.
  188. Union(inIvt IntervalTree, ivl Interval)
  189. }
  190. // NewIntervalTree returns a new interval tree.
  191. func NewIntervalTree() IntervalTree {
  192. sentinel := &intervalNode{
  193. iv: IntervalValue{},
  194. max: nil,
  195. left: nil,
  196. right: nil,
  197. parent: nil,
  198. c: black,
  199. }
  200. return &intervalTree{
  201. root: sentinel,
  202. count: 0,
  203. sentinel: sentinel,
  204. }
  205. }
  206. type intervalTree struct {
  207. root *intervalNode
  208. count int
  209. // red-black NIL node
  210. // use 'sentinel' as a dummy object to simplify boundary conditions
  211. // use the sentinel to treat a nil child of a node x as an ordinary node whose parent is x
  212. // use one shared sentinel to represent all nil leaves and the root's parent
  213. sentinel *intervalNode
  214. }
  215. // TODO: make this consistent with textbook implementation
  216. //
  217. // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.4, p324
  218. //
  219. // 0. RB-DELETE(T, z)
  220. // 1.
  221. // 2. y = z
  222. // 3. y-original-color = y.color
  223. // 4.
  224. // 5. if z.left == T.nil
  225. // 6. x = z.right
  226. // 7. RB-TRANSPLANT(T, z, z.right)
  227. // 8. else if z.right == T.nil
  228. // 9. x = z.left
  229. // 10. RB-TRANSPLANT(T, z, z.left)
  230. // 11. else
  231. // 12. y = TREE-MINIMUM(z.right)
  232. // 13. y-original-color = y.color
  233. // 14. x = y.right
  234. // 15. if y.p == z
  235. // 16. x.p = y
  236. // 17. else
  237. // 18. RB-TRANSPLANT(T, y, y.right)
  238. // 19. y.right = z.right
  239. // 20. y.right.p = y
  240. // 21. RB-TRANSPLANT(T, z, y)
  241. // 22. y.left = z.left
  242. // 23. y.left.p = y
  243. // 24. y.color = z.color
  244. // 25.
  245. // 26. if y-original-color == BLACK
  246. // 27. RB-DELETE-FIXUP(T, x)
  247. // Delete removes the node with the given interval from the tree, returning
  248. // true if a node is in fact removed.
  249. func (ivt *intervalTree) Delete(ivl Interval) bool {
  250. z := ivt.find(ivl)
  251. if z == ivt.sentinel {
  252. return false
  253. }
  254. y := z
  255. if z.left != ivt.sentinel && z.right != ivt.sentinel {
  256. y = z.successor(ivt.sentinel)
  257. }
  258. x := ivt.sentinel
  259. if y.left != ivt.sentinel {
  260. x = y.left
  261. } else if y.right != ivt.sentinel {
  262. x = y.right
  263. }
  264. x.parent = y.parent
  265. if y.parent == ivt.sentinel {
  266. ivt.root = x
  267. } else {
  268. if y == y.parent.left {
  269. y.parent.left = x
  270. } else {
  271. y.parent.right = x
  272. }
  273. y.parent.updateMax(ivt.sentinel)
  274. }
  275. if y != z {
  276. z.iv = y.iv
  277. z.updateMax(ivt.sentinel)
  278. }
  279. if y.color(ivt.sentinel) == black {
  280. ivt.deleteFixup(x)
  281. }
  282. ivt.count--
  283. return true
  284. }
  285. // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.4, p326
  286. //
  287. // 0. RB-DELETE-FIXUP(T, z)
  288. // 1.
  289. // 2. while x ≠ T.root and x.color == BLACK
  290. // 3. if x == x.p.left
  291. // 4. w = x.p.right
  292. // 5. if w.color == RED
  293. // 6. w.color = BLACK
  294. // 7. x.p.color = RED
  295. // 8. LEFT-ROTATE(T, x, p)
  296. // 9. if w.left.color == BLACK and w.right.color == BLACK
  297. // 10. w.color = RED
  298. // 11. x = x.p
  299. // 12. else if w.right.color == BLACK
  300. // 13. w.left.color = BLACK
  301. // 14. w.color = RED
  302. // 15. RIGHT-ROTATE(T, w)
  303. // 16. w = w.p.right
  304. // 17. w.color = x.p.color
  305. // 18. x.p.color = BLACK
  306. // 19. LEFT-ROTATE(T, w.p)
  307. // 20. x = T.root
  308. // 21. else
  309. // 22. w = x.p.left
  310. // 23. if w.color == RED
  311. // 24. w.color = BLACK
  312. // 25. x.p.color = RED
  313. // 26. RIGHT-ROTATE(T, x, p)
  314. // 27. if w.right.color == BLACK and w.left.color == BLACK
  315. // 28. w.color = RED
  316. // 29. x = x.p
  317. // 30. else if w.left.color == BLACK
  318. // 31. w.right.color = BLACK
  319. // 32. w.color = RED
  320. // 33. LEFT-ROTATE(T, w)
  321. // 34. w = w.p.left
  322. // 35. w.color = x.p.color
  323. // 36. x.p.color = BLACK
  324. // 37. RIGHT-ROTATE(T, w.p)
  325. // 38. x = T.root
  326. // 39.
  327. // 40. x.color = BLACK
  328. //
  329. func (ivt *intervalTree) deleteFixup(x *intervalNode) {
  330. for x != ivt.root && x.color(ivt.sentinel) == black {
  331. if x == x.parent.left { // line 3-20
  332. w := x.parent.right
  333. if w.color(ivt.sentinel) == red {
  334. w.c = black
  335. x.parent.c = red
  336. ivt.rotateLeft(x.parent)
  337. w = x.parent.right
  338. }
  339. if w == nil {
  340. break
  341. }
  342. if w.left.color(ivt.sentinel) == black && w.right.color(ivt.sentinel) == black {
  343. w.c = red
  344. x = x.parent
  345. } else {
  346. if w.right.color(ivt.sentinel) == black {
  347. w.left.c = black
  348. w.c = red
  349. ivt.rotateRight(w)
  350. w = x.parent.right
  351. }
  352. w.c = x.parent.color(ivt.sentinel)
  353. x.parent.c = black
  354. w.right.c = black
  355. ivt.rotateLeft(x.parent)
  356. x = ivt.root
  357. }
  358. } else { // line 22-38
  359. // same as above but with left and right exchanged
  360. w := x.parent.left
  361. if w.color(ivt.sentinel) == red {
  362. w.c = black
  363. x.parent.c = red
  364. ivt.rotateRight(x.parent)
  365. w = x.parent.left
  366. }
  367. if w == nil {
  368. break
  369. }
  370. if w.left.color(ivt.sentinel) == black && w.right.color(ivt.sentinel) == black {
  371. w.c = red
  372. x = x.parent
  373. } else {
  374. if w.left.color(ivt.sentinel) == black {
  375. w.right.c = black
  376. w.c = red
  377. ivt.rotateLeft(w)
  378. w = x.parent.left
  379. }
  380. w.c = x.parent.color(ivt.sentinel)
  381. x.parent.c = black
  382. w.left.c = black
  383. ivt.rotateRight(x.parent)
  384. x = ivt.root
  385. }
  386. }
  387. }
  388. if x != nil {
  389. x.c = black
  390. }
  391. }
  392. func (ivt *intervalTree) createIntervalNode(ivl Interval, val interface{}) *intervalNode {
  393. return &intervalNode{
  394. iv: IntervalValue{ivl, val},
  395. max: ivl.End,
  396. c: red,
  397. left: ivt.sentinel,
  398. right: ivt.sentinel,
  399. parent: ivt.sentinel,
  400. }
  401. }
  402. // TODO: make this consistent with textbook implementation
  403. //
  404. // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.3, p315
  405. //
  406. // 0. RB-INSERT(T, z)
  407. // 1.
  408. // 2. y = T.nil
  409. // 3. x = T.root
  410. // 4.
  411. // 5. while x ≠ T.nil
  412. // 6. y = x
  413. // 7. if z.key < x.key
  414. // 8. x = x.left
  415. // 9. else
  416. // 10. x = x.right
  417. // 11.
  418. // 12. z.p = y
  419. // 13.
  420. // 14. if y == T.nil
  421. // 15. T.root = z
  422. // 16. else if z.key < y.key
  423. // 17. y.left = z
  424. // 18. else
  425. // 19. y.right = z
  426. // 20.
  427. // 21. z.left = T.nil
  428. // 22. z.right = T.nil
  429. // 23. z.color = RED
  430. // 24.
  431. // 25. RB-INSERT-FIXUP(T, z)
  432. // Insert adds a node with the given interval into the tree.
  433. func (ivt *intervalTree) Insert(ivl Interval, val interface{}) {
  434. y := ivt.sentinel
  435. z := ivt.createIntervalNode(ivl, val)
  436. x := ivt.root
  437. for x != ivt.sentinel {
  438. y = x
  439. if z.iv.Ivl.Begin.Compare(x.iv.Ivl.Begin) < 0 {
  440. x = x.left
  441. } else {
  442. x = x.right
  443. }
  444. }
  445. z.parent = y
  446. if y == ivt.sentinel {
  447. ivt.root = z
  448. } else {
  449. if z.iv.Ivl.Begin.Compare(y.iv.Ivl.Begin) < 0 {
  450. y.left = z
  451. } else {
  452. y.right = z
  453. }
  454. y.updateMax(ivt.sentinel)
  455. }
  456. z.c = red
  457. ivt.insertFixup(z)
  458. ivt.count++
  459. }
  460. // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.3, p316
  461. //
  462. // 0. RB-INSERT-FIXUP(T, z)
  463. // 1.
  464. // 2. while z.p.color == RED
  465. // 3. if z.p == z.p.p.left
  466. // 4. y = z.p.p.right
  467. // 5. if y.color == RED
  468. // 6. z.p.color = BLACK
  469. // 7. y.color = BLACK
  470. // 8. z.p.p.color = RED
  471. // 9. z = z.p.p
  472. // 10. else if z == z.p.right
  473. // 11. z = z.p
  474. // 12. LEFT-ROTATE(T, z)
  475. // 13. z.p.color = BLACK
  476. // 14. z.p.p.color = RED
  477. // 15. RIGHT-ROTATE(T, z.p.p)
  478. // 16. else
  479. // 17. y = z.p.p.left
  480. // 18. if y.color == RED
  481. // 19. z.p.color = BLACK
  482. // 20. y.color = BLACK
  483. // 21. z.p.p.color = RED
  484. // 22. z = z.p.p
  485. // 23. else if z == z.p.right
  486. // 24. z = z.p
  487. // 25. RIGHT-ROTATE(T, z)
  488. // 26. z.p.color = BLACK
  489. // 27. z.p.p.color = RED
  490. // 28. LEFT-ROTATE(T, z.p.p)
  491. // 29.
  492. // 30. T.root.color = BLACK
  493. //
  494. func (ivt *intervalTree) insertFixup(z *intervalNode) {
  495. for z.parent.color(ivt.sentinel) == red {
  496. if z.parent == z.parent.parent.left { // line 3-15
  497. y := z.parent.parent.right
  498. if y.color(ivt.sentinel) == red {
  499. y.c = black
  500. z.parent.c = black
  501. z.parent.parent.c = red
  502. z = z.parent.parent
  503. } else {
  504. if z == z.parent.right {
  505. z = z.parent
  506. ivt.rotateLeft(z)
  507. }
  508. z.parent.c = black
  509. z.parent.parent.c = red
  510. ivt.rotateRight(z.parent.parent)
  511. }
  512. } else { // line 16-28
  513. // same as then with left/right exchanged
  514. y := z.parent.parent.left
  515. if y.color(ivt.sentinel) == red {
  516. y.c = black
  517. z.parent.c = black
  518. z.parent.parent.c = red
  519. z = z.parent.parent
  520. } else {
  521. if z == z.parent.left {
  522. z = z.parent
  523. ivt.rotateRight(z)
  524. }
  525. z.parent.c = black
  526. z.parent.parent.c = red
  527. ivt.rotateLeft(z.parent.parent)
  528. }
  529. }
  530. }
  531. // line 30
  532. ivt.root.c = black
  533. }
  534. // rotateLeft moves x so it is left of its right child
  535. //
  536. // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.2, p313
  537. //
  538. // 0. LEFT-ROTATE(T, x)
  539. // 1.
  540. // 2. y = x.right
  541. // 3. x.right = y.left
  542. // 4.
  543. // 5. if y.left ≠ T.nil
  544. // 6. y.left.p = x
  545. // 7.
  546. // 8. y.p = x.p
  547. // 9.
  548. // 10. if x.p == T.nil
  549. // 11. T.root = y
  550. // 12. else if x == x.p.left
  551. // 13. x.p.left = y
  552. // 14. else
  553. // 15. x.p.right = y
  554. // 16.
  555. // 17. y.left = x
  556. // 18. x.p = y
  557. //
  558. func (ivt *intervalTree) rotateLeft(x *intervalNode) {
  559. // rotateLeft x must have right child
  560. if x.right == ivt.sentinel {
  561. return
  562. }
  563. // line 2-3
  564. y := x.right
  565. x.right = y.left
  566. // line 5-6
  567. if y.left != ivt.sentinel {
  568. y.left.parent = x
  569. }
  570. x.updateMax(ivt.sentinel)
  571. // line 10-15, 18
  572. ivt.replaceParent(x, y)
  573. // line 17
  574. y.left = x
  575. y.updateMax(ivt.sentinel)
  576. }
  577. // rotateRight moves x so it is right of its left child
  578. //
  579. // 0. RIGHT-ROTATE(T, x)
  580. // 1.
  581. // 2. y = x.left
  582. // 3. x.left = y.right
  583. // 4.
  584. // 5. if y.right ≠ T.nil
  585. // 6. y.right.p = x
  586. // 7.
  587. // 8. y.p = x.p
  588. // 9.
  589. // 10. if x.p == T.nil
  590. // 11. T.root = y
  591. // 12. else if x == x.p.right
  592. // 13. x.p.right = y
  593. // 14. else
  594. // 15. x.p.left = y
  595. // 16.
  596. // 17. y.right = x
  597. // 18. x.p = y
  598. //
  599. func (ivt *intervalTree) rotateRight(x *intervalNode) {
  600. // rotateRight x must have left child
  601. if x.left == ivt.sentinel {
  602. return
  603. }
  604. // line 2-3
  605. y := x.left
  606. x.left = y.right
  607. // line 5-6
  608. if y.right != ivt.sentinel {
  609. y.right.parent = x
  610. }
  611. x.updateMax(ivt.sentinel)
  612. // line 10-15, 18
  613. ivt.replaceParent(x, y)
  614. // line 17
  615. y.right = x
  616. y.updateMax(ivt.sentinel)
  617. }
  618. // replaceParent replaces x's parent with y
  619. func (ivt *intervalTree) replaceParent(x *intervalNode, y *intervalNode) {
  620. y.parent = x.parent
  621. if x.parent == ivt.sentinel {
  622. ivt.root = y
  623. } else {
  624. if x == x.parent.left {
  625. x.parent.left = y
  626. } else {
  627. x.parent.right = y
  628. }
  629. x.parent.updateMax(ivt.sentinel)
  630. }
  631. x.parent = y
  632. }
  633. // Len gives the number of elements in the tree
  634. func (ivt *intervalTree) Len() int { return ivt.count }
  635. // Height is the number of levels in the tree; one node has height 1.
  636. func (ivt *intervalTree) Height() int { return ivt.root.height(ivt.sentinel) }
  637. // MaxHeight is the expected maximum tree height given the number of nodes
  638. func (ivt *intervalTree) MaxHeight() int {
  639. return int((2 * math.Log2(float64(ivt.Len()+1))) + 0.5)
  640. }
  641. // IntervalVisitor is used on tree searches; return false to stop searching.
  642. type IntervalVisitor func(n *IntervalValue) bool
  643. // Visit calls a visitor function on every tree node intersecting the given interval.
  644. // It will visit each interval [x, y) in ascending order sorted on x.
  645. func (ivt *intervalTree) Visit(ivl Interval, ivv IntervalVisitor) {
  646. ivt.root.visit(&ivl, ivt.sentinel, func(n *intervalNode) bool { return ivv(&n.iv) })
  647. }
  648. // find the exact node for a given interval
  649. func (ivt *intervalTree) find(ivl Interval) *intervalNode {
  650. ret := ivt.sentinel
  651. f := func(n *intervalNode) bool {
  652. if n.iv.Ivl != ivl {
  653. return true
  654. }
  655. ret = n
  656. return false
  657. }
  658. ivt.root.visit(&ivl, ivt.sentinel, f)
  659. return ret
  660. }
  661. // Find gets the IntervalValue for the node matching the given interval
  662. func (ivt *intervalTree) Find(ivl Interval) (ret *IntervalValue) {
  663. n := ivt.find(ivl)
  664. if n == ivt.sentinel {
  665. return nil
  666. }
  667. return &n.iv
  668. }
  669. // Intersects returns true if there is some tree node intersecting the given interval.
  670. func (ivt *intervalTree) Intersects(iv Interval) bool {
  671. x := ivt.root
  672. for x != ivt.sentinel && iv.Compare(&x.iv.Ivl) != 0 {
  673. if x.left != ivt.sentinel && x.left.max.Compare(iv.Begin) > 0 {
  674. x = x.left
  675. } else {
  676. x = x.right
  677. }
  678. }
  679. return x != ivt.sentinel
  680. }
  681. // Contains returns true if the interval tree's keys cover the entire given interval.
  682. func (ivt *intervalTree) Contains(ivl Interval) bool {
  683. var maxEnd, minBegin Comparable
  684. isContiguous := true
  685. ivt.Visit(ivl, func(n *IntervalValue) bool {
  686. if minBegin == nil {
  687. minBegin = n.Ivl.Begin
  688. maxEnd = n.Ivl.End
  689. return true
  690. }
  691. if maxEnd.Compare(n.Ivl.Begin) < 0 {
  692. isContiguous = false
  693. return false
  694. }
  695. if n.Ivl.End.Compare(maxEnd) > 0 {
  696. maxEnd = n.Ivl.End
  697. }
  698. return true
  699. })
  700. return isContiguous && minBegin != nil && maxEnd.Compare(ivl.End) >= 0 && minBegin.Compare(ivl.Begin) <= 0
  701. }
  702. // Stab returns a slice with all elements in the tree intersecting the interval.
  703. func (ivt *intervalTree) Stab(iv Interval) (ivs []*IntervalValue) {
  704. if ivt.count == 0 {
  705. return nil
  706. }
  707. f := func(n *IntervalValue) bool { ivs = append(ivs, n); return true }
  708. ivt.Visit(iv, f)
  709. return ivs
  710. }
  711. // Union merges a given interval tree into the receiver.
  712. func (ivt *intervalTree) Union(inIvt IntervalTree, ivl Interval) {
  713. f := func(n *IntervalValue) bool {
  714. ivt.Insert(n.Ivl, n.Val)
  715. return true
  716. }
  717. inIvt.Visit(ivl, f)
  718. }
  719. type visitedInterval struct {
  720. root Interval
  721. left Interval
  722. right Interval
  723. color rbcolor
  724. depth int
  725. }
  726. func (vi visitedInterval) String() string {
  727. bd := new(strings.Builder)
  728. bd.WriteString(fmt.Sprintf("root [%v,%v,%v], left [%v,%v], right [%v,%v], depth %d",
  729. vi.root.Begin, vi.root.End, vi.color,
  730. vi.left.Begin, vi.left.End,
  731. vi.right.Begin, vi.right.End,
  732. vi.depth,
  733. ))
  734. return bd.String()
  735. }
  736. // visitLevel traverses tree in level order.
  737. // used for testing
  738. func (ivt *intervalTree) visitLevel() []visitedInterval {
  739. if ivt.root == ivt.sentinel {
  740. return nil
  741. }
  742. rs := make([]visitedInterval, 0, ivt.Len())
  743. type pair struct {
  744. node *intervalNode
  745. depth int
  746. }
  747. queue := []pair{{ivt.root, 0}}
  748. for len(queue) > 0 {
  749. f := queue[0]
  750. queue = queue[1:]
  751. vi := visitedInterval{
  752. root: f.node.iv.Ivl,
  753. color: f.node.color(ivt.sentinel),
  754. depth: f.depth,
  755. }
  756. if f.node.left != ivt.sentinel {
  757. vi.left = f.node.left.iv.Ivl
  758. queue = append(queue, pair{f.node.left, f.depth + 1})
  759. }
  760. if f.node.right != ivt.sentinel {
  761. vi.right = f.node.right.iv.Ivl
  762. queue = append(queue, pair{f.node.right, f.depth + 1})
  763. }
  764. rs = append(rs, vi)
  765. }
  766. return rs
  767. }
  768. type StringComparable string
  769. func (s StringComparable) Compare(c Comparable) int {
  770. sc := c.(StringComparable)
  771. if s < sc {
  772. return -1
  773. }
  774. if s > sc {
  775. return 1
  776. }
  777. return 0
  778. }
  779. func NewStringInterval(begin, end string) Interval {
  780. return Interval{StringComparable(begin), StringComparable(end)}
  781. }
  782. func NewStringPoint(s string) Interval {
  783. return Interval{StringComparable(s), StringComparable(s + "\x00")}
  784. }
  785. // StringAffineComparable treats "" as > all other strings
  786. type StringAffineComparable string
  787. func (s StringAffineComparable) Compare(c Comparable) int {
  788. sc := c.(StringAffineComparable)
  789. if len(s) == 0 {
  790. if len(sc) == 0 {
  791. return 0
  792. }
  793. return 1
  794. }
  795. if len(sc) == 0 {
  796. return -1
  797. }
  798. if s < sc {
  799. return -1
  800. }
  801. if s > sc {
  802. return 1
  803. }
  804. return 0
  805. }
  806. func NewStringAffineInterval(begin, end string) Interval {
  807. return Interval{StringAffineComparable(begin), StringAffineComparable(end)}
  808. }
  809. func NewStringAffinePoint(s string) Interval {
  810. return NewStringAffineInterval(s, s+"\x00")
  811. }
  812. func NewInt64Interval(a int64, b int64) Interval {
  813. return Interval{Int64Comparable(a), Int64Comparable(b)}
  814. }
  815. func newInt64EmptyInterval() Interval {
  816. return Interval{Begin: nil, End: nil}
  817. }
  818. func NewInt64Point(a int64) Interval {
  819. return Interval{Int64Comparable(a), Int64Comparable(a + 1)}
  820. }
  821. type Int64Comparable int64
  822. func (v Int64Comparable) Compare(c Comparable) int {
  823. vc := c.(Int64Comparable)
  824. cmp := v - vc
  825. if cmp < 0 {
  826. return -1
  827. }
  828. if cmp > 0 {
  829. return 1
  830. }
  831. return 0
  832. }
  833. // BytesAffineComparable treats empty byte arrays as > all other byte arrays
  834. type BytesAffineComparable []byte
  835. func (b BytesAffineComparable) Compare(c Comparable) int {
  836. bc := c.(BytesAffineComparable)
  837. if len(b) == 0 {
  838. if len(bc) == 0 {
  839. return 0
  840. }
  841. return 1
  842. }
  843. if len(bc) == 0 {
  844. return -1
  845. }
  846. return bytes.Compare(b, bc)
  847. }
  848. func NewBytesAffineInterval(begin, end []byte) Interval {
  849. return Interval{BytesAffineComparable(begin), BytesAffineComparable(end)}
  850. }
  851. func NewBytesAffinePoint(b []byte) Interval {
  852. be := make([]byte, len(b)+1)
  853. copy(be, b)
  854. be[len(b)] = 0
  855. return NewBytesAffineInterval(b, be)
  856. }