interval_tree.go 20 KB

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