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