Matthias Andreas Benkard | 832a54e | 2019-01-29 09:27:38 +0100 | [diff] [blame] | 1 | // Copyright 2010 Petar Maymounkov. All rights reserved. |
| 2 | // Use of this source code is governed by a BSD-style |
| 3 | // license that can be found in the LICENSE file. |
| 4 | |
| 5 | // A Left-Leaning Red-Black (LLRB) implementation of 2-3 balanced binary search trees, |
| 6 | // based on the following work: |
| 7 | // |
| 8 | // http://www.cs.princeton.edu/~rs/talks/LLRB/08Penn.pdf |
| 9 | // http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf |
| 10 | // http://www.cs.princeton.edu/~rs/talks/LLRB/Java/RedBlackBST.java |
| 11 | // |
| 12 | // 2-3 trees (and the run-time equivalent 2-3-4 trees) are the de facto standard BST |
| 13 | // algoritms found in implementations of Python, Java, and other libraries. The LLRB |
| 14 | // implementation of 2-3 trees is a recent improvement on the traditional implementation, |
| 15 | // observed and documented by Robert Sedgewick. |
| 16 | // |
| 17 | package llrb |
| 18 | |
| 19 | // Tree is a Left-Leaning Red-Black (LLRB) implementation of 2-3 trees |
| 20 | type LLRB struct { |
| 21 | count int |
| 22 | root *Node |
| 23 | } |
| 24 | |
| 25 | type Node struct { |
| 26 | Item |
| 27 | Left, Right *Node // Pointers to left and right child nodes |
| 28 | Black bool // If set, the color of the link (incoming from the parent) is black |
| 29 | // In the LLRB, new nodes are always red, hence the zero-value for node |
| 30 | } |
| 31 | |
| 32 | type Item interface { |
| 33 | Less(than Item) bool |
| 34 | } |
| 35 | |
| 36 | // |
| 37 | func less(x, y Item) bool { |
| 38 | if x == pinf { |
| 39 | return false |
| 40 | } |
| 41 | if x == ninf { |
| 42 | return true |
| 43 | } |
| 44 | return x.Less(y) |
| 45 | } |
| 46 | |
| 47 | // Inf returns an Item that is "bigger than" any other item, if sign is positive. |
| 48 | // Otherwise it returns an Item that is "smaller than" any other item. |
| 49 | func Inf(sign int) Item { |
| 50 | if sign == 0 { |
| 51 | panic("sign") |
| 52 | } |
| 53 | if sign > 0 { |
| 54 | return pinf |
| 55 | } |
| 56 | return ninf |
| 57 | } |
| 58 | |
| 59 | var ( |
| 60 | ninf = nInf{} |
| 61 | pinf = pInf{} |
| 62 | ) |
| 63 | |
| 64 | type nInf struct{} |
| 65 | |
| 66 | func (nInf) Less(Item) bool { |
| 67 | return true |
| 68 | } |
| 69 | |
| 70 | type pInf struct{} |
| 71 | |
| 72 | func (pInf) Less(Item) bool { |
| 73 | return false |
| 74 | } |
| 75 | |
| 76 | // New() allocates a new tree |
| 77 | func New() *LLRB { |
| 78 | return &LLRB{} |
| 79 | } |
| 80 | |
| 81 | // SetRoot sets the root node of the tree. |
| 82 | // It is intended to be used by functions that deserialize the tree. |
| 83 | func (t *LLRB) SetRoot(r *Node) { |
| 84 | t.root = r |
| 85 | } |
| 86 | |
| 87 | // Root returns the root node of the tree. |
| 88 | // It is intended to be used by functions that serialize the tree. |
| 89 | func (t *LLRB) Root() *Node { |
| 90 | return t.root |
| 91 | } |
| 92 | |
| 93 | // Len returns the number of nodes in the tree. |
| 94 | func (t *LLRB) Len() int { return t.count } |
| 95 | |
| 96 | // Has returns true if the tree contains an element whose order is the same as that of key. |
| 97 | func (t *LLRB) Has(key Item) bool { |
| 98 | return t.Get(key) != nil |
| 99 | } |
| 100 | |
| 101 | // Get retrieves an element from the tree whose order is the same as that of key. |
| 102 | func (t *LLRB) Get(key Item) Item { |
| 103 | h := t.root |
| 104 | for h != nil { |
| 105 | switch { |
| 106 | case less(key, h.Item): |
| 107 | h = h.Left |
| 108 | case less(h.Item, key): |
| 109 | h = h.Right |
| 110 | default: |
| 111 | return h.Item |
| 112 | } |
| 113 | } |
| 114 | return nil |
| 115 | } |
| 116 | |
| 117 | // Min returns the minimum element in the tree. |
| 118 | func (t *LLRB) Min() Item { |
| 119 | h := t.root |
| 120 | if h == nil { |
| 121 | return nil |
| 122 | } |
| 123 | for h.Left != nil { |
| 124 | h = h.Left |
| 125 | } |
| 126 | return h.Item |
| 127 | } |
| 128 | |
| 129 | // Max returns the maximum element in the tree. |
| 130 | func (t *LLRB) Max() Item { |
| 131 | h := t.root |
| 132 | if h == nil { |
| 133 | return nil |
| 134 | } |
| 135 | for h.Right != nil { |
| 136 | h = h.Right |
| 137 | } |
| 138 | return h.Item |
| 139 | } |
| 140 | |
| 141 | func (t *LLRB) ReplaceOrInsertBulk(items ...Item) { |
| 142 | for _, i := range items { |
| 143 | t.ReplaceOrInsert(i) |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | func (t *LLRB) InsertNoReplaceBulk(items ...Item) { |
| 148 | for _, i := range items { |
| 149 | t.InsertNoReplace(i) |
| 150 | } |
| 151 | } |
| 152 | |
| 153 | // ReplaceOrInsert inserts item into the tree. If an existing |
| 154 | // element has the same order, it is removed from the tree and returned. |
| 155 | func (t *LLRB) ReplaceOrInsert(item Item) Item { |
| 156 | if item == nil { |
| 157 | panic("inserting nil item") |
| 158 | } |
| 159 | var replaced Item |
| 160 | t.root, replaced = t.replaceOrInsert(t.root, item) |
| 161 | t.root.Black = true |
| 162 | if replaced == nil { |
| 163 | t.count++ |
| 164 | } |
| 165 | return replaced |
| 166 | } |
| 167 | |
| 168 | func (t *LLRB) replaceOrInsert(h *Node, item Item) (*Node, Item) { |
| 169 | if h == nil { |
| 170 | return newNode(item), nil |
| 171 | } |
| 172 | |
| 173 | h = walkDownRot23(h) |
| 174 | |
| 175 | var replaced Item |
| 176 | if less(item, h.Item) { // BUG |
| 177 | h.Left, replaced = t.replaceOrInsert(h.Left, item) |
| 178 | } else if less(h.Item, item) { |
| 179 | h.Right, replaced = t.replaceOrInsert(h.Right, item) |
| 180 | } else { |
| 181 | replaced, h.Item = h.Item, item |
| 182 | } |
| 183 | |
| 184 | h = walkUpRot23(h) |
| 185 | |
| 186 | return h, replaced |
| 187 | } |
| 188 | |
| 189 | // InsertNoReplace inserts item into the tree. If an existing |
| 190 | // element has the same order, both elements remain in the tree. |
| 191 | func (t *LLRB) InsertNoReplace(item Item) { |
| 192 | if item == nil { |
| 193 | panic("inserting nil item") |
| 194 | } |
| 195 | t.root = t.insertNoReplace(t.root, item) |
| 196 | t.root.Black = true |
| 197 | t.count++ |
| 198 | } |
| 199 | |
| 200 | func (t *LLRB) insertNoReplace(h *Node, item Item) *Node { |
| 201 | if h == nil { |
| 202 | return newNode(item) |
| 203 | } |
| 204 | |
| 205 | h = walkDownRot23(h) |
| 206 | |
| 207 | if less(item, h.Item) { |
| 208 | h.Left = t.insertNoReplace(h.Left, item) |
| 209 | } else { |
| 210 | h.Right = t.insertNoReplace(h.Right, item) |
| 211 | } |
| 212 | |
| 213 | return walkUpRot23(h) |
| 214 | } |
| 215 | |
| 216 | // Rotation driver routines for 2-3 algorithm |
| 217 | |
| 218 | func walkDownRot23(h *Node) *Node { return h } |
| 219 | |
| 220 | func walkUpRot23(h *Node) *Node { |
| 221 | if isRed(h.Right) && !isRed(h.Left) { |
| 222 | h = rotateLeft(h) |
| 223 | } |
| 224 | |
| 225 | if isRed(h.Left) && isRed(h.Left.Left) { |
| 226 | h = rotateRight(h) |
| 227 | } |
| 228 | |
| 229 | if isRed(h.Left) && isRed(h.Right) { |
| 230 | flip(h) |
| 231 | } |
| 232 | |
| 233 | return h |
| 234 | } |
| 235 | |
| 236 | // Rotation driver routines for 2-3-4 algorithm |
| 237 | |
| 238 | func walkDownRot234(h *Node) *Node { |
| 239 | if isRed(h.Left) && isRed(h.Right) { |
| 240 | flip(h) |
| 241 | } |
| 242 | |
| 243 | return h |
| 244 | } |
| 245 | |
| 246 | func walkUpRot234(h *Node) *Node { |
| 247 | if isRed(h.Right) && !isRed(h.Left) { |
| 248 | h = rotateLeft(h) |
| 249 | } |
| 250 | |
| 251 | if isRed(h.Left) && isRed(h.Left.Left) { |
| 252 | h = rotateRight(h) |
| 253 | } |
| 254 | |
| 255 | return h |
| 256 | } |
| 257 | |
| 258 | // DeleteMin deletes the minimum element in the tree and returns the |
| 259 | // deleted item or nil otherwise. |
| 260 | func (t *LLRB) DeleteMin() Item { |
| 261 | var deleted Item |
| 262 | t.root, deleted = deleteMin(t.root) |
| 263 | if t.root != nil { |
| 264 | t.root.Black = true |
| 265 | } |
| 266 | if deleted != nil { |
| 267 | t.count-- |
| 268 | } |
| 269 | return deleted |
| 270 | } |
| 271 | |
| 272 | // deleteMin code for LLRB 2-3 trees |
| 273 | func deleteMin(h *Node) (*Node, Item) { |
| 274 | if h == nil { |
| 275 | return nil, nil |
| 276 | } |
| 277 | if h.Left == nil { |
| 278 | return nil, h.Item |
| 279 | } |
| 280 | |
| 281 | if !isRed(h.Left) && !isRed(h.Left.Left) { |
| 282 | h = moveRedLeft(h) |
| 283 | } |
| 284 | |
| 285 | var deleted Item |
| 286 | h.Left, deleted = deleteMin(h.Left) |
| 287 | |
| 288 | return fixUp(h), deleted |
| 289 | } |
| 290 | |
| 291 | // DeleteMax deletes the maximum element in the tree and returns |
| 292 | // the deleted item or nil otherwise |
| 293 | func (t *LLRB) DeleteMax() Item { |
| 294 | var deleted Item |
| 295 | t.root, deleted = deleteMax(t.root) |
| 296 | if t.root != nil { |
| 297 | t.root.Black = true |
| 298 | } |
| 299 | if deleted != nil { |
| 300 | t.count-- |
| 301 | } |
| 302 | return deleted |
| 303 | } |
| 304 | |
| 305 | func deleteMax(h *Node) (*Node, Item) { |
| 306 | if h == nil { |
| 307 | return nil, nil |
| 308 | } |
| 309 | if isRed(h.Left) { |
| 310 | h = rotateRight(h) |
| 311 | } |
| 312 | if h.Right == nil { |
| 313 | return nil, h.Item |
| 314 | } |
| 315 | if !isRed(h.Right) && !isRed(h.Right.Left) { |
| 316 | h = moveRedRight(h) |
| 317 | } |
| 318 | var deleted Item |
| 319 | h.Right, deleted = deleteMax(h.Right) |
| 320 | |
| 321 | return fixUp(h), deleted |
| 322 | } |
| 323 | |
| 324 | // Delete deletes an item from the tree whose key equals key. |
| 325 | // The deleted item is return, otherwise nil is returned. |
| 326 | func (t *LLRB) Delete(key Item) Item { |
| 327 | var deleted Item |
| 328 | t.root, deleted = t.delete(t.root, key) |
| 329 | if t.root != nil { |
| 330 | t.root.Black = true |
| 331 | } |
| 332 | if deleted != nil { |
| 333 | t.count-- |
| 334 | } |
| 335 | return deleted |
| 336 | } |
| 337 | |
| 338 | func (t *LLRB) delete(h *Node, item Item) (*Node, Item) { |
| 339 | var deleted Item |
| 340 | if h == nil { |
| 341 | return nil, nil |
| 342 | } |
| 343 | if less(item, h.Item) { |
| 344 | if h.Left == nil { // item not present. Nothing to delete |
| 345 | return h, nil |
| 346 | } |
| 347 | if !isRed(h.Left) && !isRed(h.Left.Left) { |
| 348 | h = moveRedLeft(h) |
| 349 | } |
| 350 | h.Left, deleted = t.delete(h.Left, item) |
| 351 | } else { |
| 352 | if isRed(h.Left) { |
| 353 | h = rotateRight(h) |
| 354 | } |
| 355 | // If @item equals @h.Item and no right children at @h |
| 356 | if !less(h.Item, item) && h.Right == nil { |
| 357 | return nil, h.Item |
| 358 | } |
| 359 | // PETAR: Added 'h.Right != nil' below |
| 360 | if h.Right != nil && !isRed(h.Right) && !isRed(h.Right.Left) { |
| 361 | h = moveRedRight(h) |
| 362 | } |
| 363 | // If @item equals @h.Item, and (from above) 'h.Right != nil' |
| 364 | if !less(h.Item, item) { |
| 365 | var subDeleted Item |
| 366 | h.Right, subDeleted = deleteMin(h.Right) |
| 367 | if subDeleted == nil { |
| 368 | panic("logic") |
| 369 | } |
| 370 | deleted, h.Item = h.Item, subDeleted |
| 371 | } else { // Else, @item is bigger than @h.Item |
| 372 | h.Right, deleted = t.delete(h.Right, item) |
| 373 | } |
| 374 | } |
| 375 | |
| 376 | return fixUp(h), deleted |
| 377 | } |
| 378 | |
| 379 | // Internal node manipulation routines |
| 380 | |
| 381 | func newNode(item Item) *Node { return &Node{Item: item} } |
| 382 | |
| 383 | func isRed(h *Node) bool { |
| 384 | if h == nil { |
| 385 | return false |
| 386 | } |
| 387 | return !h.Black |
| 388 | } |
| 389 | |
| 390 | func rotateLeft(h *Node) *Node { |
| 391 | x := h.Right |
| 392 | if x.Black { |
| 393 | panic("rotating a black link") |
| 394 | } |
| 395 | h.Right = x.Left |
| 396 | x.Left = h |
| 397 | x.Black = h.Black |
| 398 | h.Black = false |
| 399 | return x |
| 400 | } |
| 401 | |
| 402 | func rotateRight(h *Node) *Node { |
| 403 | x := h.Left |
| 404 | if x.Black { |
| 405 | panic("rotating a black link") |
| 406 | } |
| 407 | h.Left = x.Right |
| 408 | x.Right = h |
| 409 | x.Black = h.Black |
| 410 | h.Black = false |
| 411 | return x |
| 412 | } |
| 413 | |
| 414 | // REQUIRE: Left and Right children must be present |
| 415 | func flip(h *Node) { |
| 416 | h.Black = !h.Black |
| 417 | h.Left.Black = !h.Left.Black |
| 418 | h.Right.Black = !h.Right.Black |
| 419 | } |
| 420 | |
| 421 | // REQUIRE: Left and Right children must be present |
| 422 | func moveRedLeft(h *Node) *Node { |
| 423 | flip(h) |
| 424 | if isRed(h.Right.Left) { |
| 425 | h.Right = rotateRight(h.Right) |
| 426 | h = rotateLeft(h) |
| 427 | flip(h) |
| 428 | } |
| 429 | return h |
| 430 | } |
| 431 | |
| 432 | // REQUIRE: Left and Right children must be present |
| 433 | func moveRedRight(h *Node) *Node { |
| 434 | flip(h) |
| 435 | if isRed(h.Left.Left) { |
| 436 | h = rotateRight(h) |
| 437 | flip(h) |
| 438 | } |
| 439 | return h |
| 440 | } |
| 441 | |
| 442 | func fixUp(h *Node) *Node { |
| 443 | if isRed(h.Right) { |
| 444 | h = rotateLeft(h) |
| 445 | } |
| 446 | |
| 447 | if isRed(h.Left) && isRed(h.Left.Left) { |
| 448 | h = rotateRight(h) |
| 449 | } |
| 450 | |
| 451 | if isRed(h.Left) && isRed(h.Right) { |
| 452 | flip(h) |
| 453 | } |
| 454 | |
| 455 | return h |
| 456 | } |