├── .gitignore ├── AUTHORS ├── LICENSE ├── README.md ├── doc ├── Sedgewick-LLRB.pdf ├── Sedgewick-RedBlackBST.java └── Sedgewick-Talk-Penn2008.pdf ├── example └── ex1.go ├── go.mod └── llrb ├── avgvar.go ├── iterator.go ├── iterator_test.go ├── llrb-stats.go ├── llrb.go ├── llrb_test.go └── util.go /.gitignore: -------------------------------------------------------------------------------- 1 | syntax:glob 2 | *.[568ao] 3 | *.ao 4 | *.so 5 | *.pyc 6 | *.swp 7 | *.swo 8 | ._* 9 | .nfs.* 10 | [568a].out 11 | *~ 12 | *.orig 13 | *.pb.go 14 | core 15 | _obj 16 | _test 17 | src/pkg/Make.deps 18 | _testmain.go 19 | 20 | syntax:regexp 21 | ^pkg/ 22 | ^src/cmd/(.*)/6?\1$ 23 | ^.*/core.[0-9]*$ 24 | -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- 1 | Petar Maymounkov 2 | Vadim Vygonets 3 | Ian Smith 4 | Martin Bruse 5 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2010, Petar Maymounkov 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without modification, 5 | are permitted provided that the following conditions are met: 6 | 7 | (*) Redistributions of source code must retain the above copyright notice, this list 8 | of conditions and the following disclaimer. 9 | 10 | (*) Redistributions in binary form must reproduce the above copyright notice, this 11 | list of conditions and the following disclaimer in the documentation and/or 12 | other materials provided with the distribution. 13 | 14 | (*) Neither the name of Petar Maymounkov nor the names of its contributors may be 15 | used to endorse or promote products derived from this software without specific 16 | prior written permission. 17 | 18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 19 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 20 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 22 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 23 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 24 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 25 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 26 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 27 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GoLLRB 2 | 3 | GoLLRB is a Left-Leaning Red-Black (LLRB) implementation of 2-3 balanced binary 4 | search trees in Go Language. 5 | 6 | ## Overview 7 | 8 | As of this writing and to the best of the author's knowledge, 9 | Go still does not have a balanced binary search tree (BBST) data structure. 10 | These data structures are quite useful in a variety of cases. A BBST maintains 11 | elements in sorted order under dynamic updates (inserts and deletes) and can 12 | support various order-specific queries. Furthermore, in practice one often 13 | implements other common data structures like Priority Queues, using BBST's. 14 | 15 | 2-3 trees (a type of BBST's), as well as the runtime-similar 2-3-4 trees, are 16 | the de facto standard BBST algoritms found in implementations of Python, Java, 17 | and other libraries. The LLRB method of implementing 2-3 trees is a recent 18 | improvement over the traditional implementation. The LLRB approach was 19 | discovered relatively recently (in 2008) by Robert Sedgewick of Princeton 20 | University. 21 | 22 | GoLLRB is a Go implementation of LLRB 2-3 trees. 23 | 24 | ## Maturity 25 | 26 | GoLLRB has been used in some pretty heavy-weight machine learning tasks over many gigabytes of data. 27 | I consider it to be in stable, perhaps even production, shape. There are no known bugs. 28 | 29 | ## Installation 30 | 31 | With a healthy Go Language installed, simply run `go get github.com/petar/GoLLRB/llrb` 32 | 33 | ## Example 34 | 35 | package main 36 | 37 | import ( 38 | "fmt" 39 | "github.com/petar/GoLLRB/llrb" 40 | ) 41 | 42 | func lessInt(a, b interface{}) bool { return a.(int) < b.(int) } 43 | 44 | func main() { 45 | tree := llrb.New(lessInt) 46 | tree.ReplaceOrInsert(1) 47 | tree.ReplaceOrInsert(2) 48 | tree.ReplaceOrInsert(3) 49 | tree.ReplaceOrInsert(4) 50 | tree.DeleteMin() 51 | tree.Delete(4) 52 | c := tree.IterAscend() 53 | for { 54 | u := <-c 55 | if u == nil { 56 | break 57 | } 58 | fmt.Printf("%d\n", int(u.(int))) 59 | } 60 | } 61 | 62 | ## About 63 | 64 | GoLLRB was written by [Petar Maymounkov](http://pdos.csail.mit.edu/~petar/). 65 | 66 | Follow me on [Twitter @maymounkov](http://www.twitter.com/maymounkov)! 67 | -------------------------------------------------------------------------------- /doc/Sedgewick-LLRB.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/petar/GoLLRB/ae3b015fd3e9d3cdb1e0b0e3da2caf14dd3f908f/doc/Sedgewick-LLRB.pdf -------------------------------------------------------------------------------- /doc/Sedgewick-RedBlackBST.java: -------------------------------------------------------------------------------- 1 | public class RedBlackBST, Value> 2 | { 3 | private static final int BST = 0; 4 | private static final int TD234 = 1; 5 | private static final int BU23 = 2; 6 | private static final boolean RED = true; 7 | private static final boolean BLACK = false; 8 | 9 | private Node root; // root of the BST 10 | private int k; // ordinal for drawing 11 | private final int species; // species kind of tree for insert 12 | private int heightBLACK; // black height of tree 13 | 14 | RedBlackBST(int species) 15 | { this.species = species; } 16 | 17 | private class Node 18 | { 19 | Key key; // key 20 | Value value; // associated data 21 | Node left, right; // left and right subtrees 22 | boolean color; // color of parent link 23 | private int N; // number of nodes in tree rooted here 24 | private int height; // height of tree rooted here 25 | private double xc, yc; // for drawing 26 | 27 | Node(Key key, Value value) 28 | { 29 | this.key = key; 30 | this.value = value; 31 | this.color = RED; 32 | this.N = 1; 33 | this.height = 1; 34 | } 35 | } 36 | 37 | public int size() 38 | { return size(root); } 39 | 40 | private int size(Node x) 41 | { 42 | if (x == null) return 0; 43 | else return x.N; 44 | } 45 | 46 | public int rootRank() 47 | { 48 | if (root == null) return 0; 49 | else return size(root.left); 50 | } 51 | 52 | public int height() 53 | { return height(root); } 54 | 55 | public int heightB() 56 | { return heightBLACK; } 57 | 58 | private int height(Node x) 59 | { 60 | if (x == null) return 0; 61 | else return x.height; 62 | } 63 | 64 | public boolean contains(Key key) 65 | { return (get(key) != null); } 66 | 67 | public Value get(Key key) 68 | { return get(root, key); } 69 | 70 | private Value get(Node x, Key key) 71 | { 72 | if (x == null) return null; 73 | if (eq (key, x.key)) return x.value; 74 | if (less(key, x.key)) return get(x.left, key); 75 | else return get(x.right, key); 76 | } 77 | 78 | public Key min() 79 | { 80 | if (root == null) return null; 81 | else return min(root); 82 | } 83 | 84 | private Key min(Node x) 85 | { 86 | if (x.left == null) return x.key; 87 | else return min(x.left); 88 | } 89 | 90 | public Key max() 91 | { 92 | if (root == null) return null; 93 | else return max(root); 94 | } 95 | 96 | private Key max(Node x) 97 | { 98 | if (x.right == null) return x.key; 99 | else return max(x.right); 100 | } 101 | 102 | public void put(Key key, Value value) 103 | { 104 | root = insert(root, key, value); 105 | if (isRed(root)) heightBLACK++; 106 | root.color = BLACK; 107 | } 108 | 109 | private Node insert(Node h, Key key, Value value) 110 | { 111 | if (h == null) 112 | return new Node(key, value); 113 | 114 | if (species == TD234) 115 | if (isRed(h.left) && isRed(h.right)) 116 | colorFlip(h); 117 | 118 | if (eq(key, h.key)) 119 | h.value = value; 120 | else if (less(key, h.key)) 121 | h.left = insert(h.left, key, value); 122 | else 123 | h.right = insert(h.right, key, value); 124 | 125 | if (species == BST) return setN(h); 126 | 127 | if (isRed(h.right)) 128 | h = rotateLeft(h); 129 | 130 | if (isRed(h.left) && isRed(h.left.left)) 131 | h = rotateRight(h); 132 | 133 | if (species == BU23) 134 | if (isRed(h.left) && isRed(h.right)) 135 | colorFlip(h); 136 | 137 | return setN(h); 138 | } 139 | 140 | public void deleteMin() 141 | { 142 | root = deleteMin(root); 143 | root.color = BLACK; 144 | } 145 | 146 | private Node deleteMin(Node h) 147 | { 148 | if (h.left == null) 149 | return null; 150 | 151 | if (!isRed(h.left) && !isRed(h.left.left)) 152 | h = moveRedLeft(h); 153 | 154 | h.left = deleteMin(h.left); 155 | 156 | return fixUp(h); 157 | } 158 | 159 | public void deleteMax() 160 | { 161 | root = deleteMax(root); 162 | root.color = BLACK; 163 | } 164 | 165 | private Node deleteMax(Node h) 166 | { 167 | // if (h.right == null) 168 | // { 169 | // if (h.left != null) 170 | // h.left.color = BLACK; 171 | // return h.left; 172 | // } 173 | 174 | if (isRed(h.left)) 175 | h = rotateRight(h); 176 | 177 | if (h.right == null) 178 | return null; 179 | 180 | if (!isRed(h.right) && !isRed(h.right.left)) 181 | h = moveRedRight(h); 182 | 183 | h.right = deleteMax(h.right); 184 | 185 | return fixUp(h); 186 | } 187 | 188 | public void delete(Key key) 189 | { 190 | root = delete(root, key); 191 | root.color = BLACK; 192 | } 193 | 194 | private Node delete(Node h, Key key) 195 | { 196 | if (less(key, h.key)) 197 | { 198 | if (!isRed(h.left) && !isRed(h.left.left)) 199 | h = moveRedLeft(h); 200 | h.left = delete(h.left, key); 201 | } 202 | else 203 | { 204 | if (isRed(h.left)) 205 | h = rotateRight(h); 206 | if (eq(key, h.key) && (h.right == null)) 207 | return null; 208 | if (!isRed(h.right) && !isRed(h.right.left)) 209 | h = moveRedRight(h); 210 | if (eq(key, h.key)) 211 | { 212 | h.value = get(h.right, min(h.right)); 213 | h.key = min(h.right); 214 | h.right = deleteMin(h.right); 215 | } 216 | else h.right = delete(h.right, key); 217 | } 218 | 219 | return fixUp(h); 220 | } 221 | 222 | // Helper methods 223 | 224 | private boolean less(Key a, Key b) { return a.compareTo(b) < 0; } 225 | private boolean eq (Key a, Key b) { return a.compareTo(b) == 0; } 226 | 227 | private boolean isRed(Node x) 228 | { 229 | if (x == null) return false; 230 | return (x.color == RED); 231 | } 232 | 233 | private void colorFlip(Node h) 234 | { 235 | h.color = !h.color; 236 | h.left.color = !h.left.color; 237 | h.right.color = !h.right.color; 238 | } 239 | 240 | private Node rotateLeft(Node h) 241 | { // Make a right-leaning 3-node lean to the left. 242 | Node x = h.right; 243 | h.right = x.left; 244 | x.left = setN(h); 245 | x.color = x.left.color; 246 | x.left.color = RED; 247 | return setN(x); 248 | } 249 | 250 | private Node rotateRight(Node h) 251 | { // Make a left-leaning 3-node lean to the right. 252 | Node x = h.left; 253 | h.left = x.right; 254 | x.right = setN(h); 255 | x.color = x.right.color; 256 | x.right.color = RED; 257 | return setN(x); 258 | } 259 | 260 | private Node moveRedLeft(Node h) 261 | { // Assuming that h is red and both h.left and h.left.left 262 | // are black, make h.left or one of its children red. 263 | colorFlip(h); 264 | if (isRed(h.right.left)) 265 | { 266 | h.right = rotateRight(h.right); 267 | h = rotateLeft(h); 268 | colorFlip(h); 269 | } 270 | return h; 271 | } 272 | 273 | private Node moveRedRight(Node h) 274 | { // Assuming that h is red and both h.right and h.right.left 275 | // are black, make h.right or one of its children red. 276 | colorFlip(h); 277 | if (isRed(h.left.left)) 278 | { 279 | h = rotateRight(h); 280 | colorFlip(h); 281 | } 282 | return h; 283 | } 284 | 285 | private Node fixUp(Node h) 286 | { 287 | if (isRed(h.right)) 288 | h = rotateLeft(h); 289 | 290 | if (isRed(h.left) && isRed(h.left.left)) 291 | h = rotateRight(h); 292 | 293 | if (isRed(h.left) && isRed(h.right)) 294 | colorFlip(h); 295 | 296 | return setN(h); 297 | } 298 | 299 | private Node setN(Node h) 300 | { 301 | h.N = size(h.left) + size(h.right) + 1; 302 | if (height(h.left) > height(h.right)) h.height = height(h.left) + 1; 303 | else h.height = height(h.right) + 1; 304 | return h; 305 | } 306 | 307 | public String toString() 308 | { 309 | if (root == null) return ""; 310 | else return heightB() + " " + toString(root); 311 | } 312 | 313 | public String toString(Node x) 314 | { 315 | String s = "("; 316 | if (x.left == null) s += "("; else s += toString(x.left); 317 | if (isRed(x)) s += "*"; 318 | if (x.right == null) s += ")"; else s += toString(x.right); 319 | return s + ")"; 320 | } 321 | 322 | // Methods for tree drawing 323 | 324 | public void draw(double y, double lineWidth, double nodeSize) 325 | { 326 | k = 0; 327 | setcoords(root, y); 328 | StdDraw.setPenColor(StdDraw.BLACK); 329 | StdDraw.setPenRadius(lineWidth); 330 | drawlines(root); 331 | StdDraw.setPenColor(StdDraw.WHITE); 332 | drawnodes(root, nodeSize); 333 | } 334 | 335 | public void setcoords(Node x, double d) 336 | { 337 | if (x == null) return; 338 | setcoords(x.left, d-.04); 339 | x.xc = (0.5 + k++)/size(); x.yc = d - .04; 340 | setcoords(x.right, d-.04); 341 | } 342 | 343 | public void drawlines(Node x) 344 | { 345 | if (x == null) return; 346 | drawlines(x.left); 347 | if (x.left != null) 348 | { 349 | if (x.left.color == RED) StdDraw.setPenColor(StdDraw.RED); 350 | else StdDraw.setPenColor(StdDraw.BLACK); 351 | StdDraw.line(x.xc, x.yc, x.left.xc, x.left.yc); 352 | } 353 | if (x.right != null) 354 | { 355 | if (x.right.color == RED) StdDraw.setPenColor(StdDraw.RED); 356 | else StdDraw.setPenColor(StdDraw.BLACK); 357 | StdDraw.line(x.xc, x.yc, x.right.xc, x.right.yc); 358 | } 359 | drawlines(x.right); 360 | } 361 | 362 | public void drawnodes(Node x, double nodeSize) 363 | { 364 | if (x == null) return; 365 | drawnodes(x.left, nodeSize); 366 | StdDraw.filledCircle(x.xc, x.yc, nodeSize); 367 | drawnodes(x.right, nodeSize); 368 | } 369 | 370 | public void mark(Key key) 371 | { 372 | StdDraw.setPenColor(StdDraw.BLACK); 373 | marknodes(key, root); 374 | } 375 | 376 | public void marknodes(Key key, Node x) 377 | { 378 | if (x == null) return; 379 | marknodes(key, x.left); 380 | if (eq(key, x.key)) 381 | StdDraw.filledCircle(x.xc, x.yc, .004); 382 | marknodes(key, x.right); 383 | } 384 | 385 | public int ipl() 386 | { return ipl(root); } 387 | 388 | public int ipl(Node x) 389 | { 390 | if (x == null) return 0; 391 | return size(x) - 1 + ipl(x.left) + ipl(x.right); 392 | } 393 | 394 | public int sizeRed() 395 | { return sizeRed(root); } 396 | 397 | public int sizeRed(Node x) 398 | { 399 | if (x == null) return 0; 400 | if (isRed(x)) return 1 + sizeRed(x.left) + sizeRed(x.right); 401 | else return sizeRed(x.left) + sizeRed(x.right); 402 | } 403 | 404 | // Integrity checks 405 | 406 | public boolean check() 407 | { // Is this tree a red-black tree? 408 | return isBST() && is234() && isBalanced(); 409 | } 410 | 411 | private boolean isBST() 412 | { // Is this tree a BST? 413 | return isBST(root, min(), max()); 414 | } 415 | 416 | private boolean isBST(Node x, Key min, Key max) 417 | { // Are all the values in the BST rooted at x between min and max, 418 | // and does the same property hold for both subtrees? 419 | if (x == null) return true; 420 | if (less(x.key, min) || less(max, x.key)) return false; 421 | return isBST(x.left, min, x.key) && isBST(x.right, x.key, max); 422 | } 423 | 424 | private boolean is234() { return is234(root); } 425 | private boolean is234(Node x) 426 | { // Does the tree have no red right links, and at most two (left) 427 | // red links in a row on any path? 428 | if (x == null) return true; 429 | if (isRed(x.right)) return false; 430 | if (isRed(x)) 431 | if (isRed(x.left)) 432 | if (isRed(x.left.left)) return false; 433 | return is234(x.left) && is234(x.right); 434 | } 435 | 436 | private boolean isBalanced() 437 | { // Do all paths from root to leaf have same number of black edges? 438 | int black = 0; // number of black links on path from root to min 439 | Node x = root; 440 | while (x != null) 441 | { 442 | if (!isRed(x)) black++; 443 | x = x.left; 444 | } 445 | return isBalanced(root, black); 446 | } 447 | 448 | private boolean isBalanced(Node x, int black) 449 | { // Does every path from the root to a leaf have the given number 450 | // of black links? 451 | if (x == null && black == 0) return true; 452 | else if (x == null && black != 0) return false; 453 | if (!isRed(x)) black--; 454 | return isBalanced(x.left, black) && isBalanced(x.right, black); 455 | } 456 | 457 | 458 | public static void main(String[] args) 459 | { 460 | StdDraw.setPenRadius(.0025); 461 | int species = Integer.parseInt(args[0]); 462 | RedBlackBST st; 463 | st = new RedBlackBST(species); 464 | int[] a = { 3, 1, 4, 2, 5, 9, 6, 8, 7 }; 465 | for (int i = 0; i < a.length; i++) 466 | st.put(a[i], i); 467 | StdOut.println(st); 468 | StdDraw.clear(StdDraw.LIGHT_GRAY); 469 | st.draw(.95, .0025, .008); 470 | StdOut.println(st.min() + " " + st.max() + " " + st.check()); 471 | StdOut.println(st.ipl()); 472 | StdOut.println(st.heightB()); 473 | } 474 | 475 | } 476 | -------------------------------------------------------------------------------- /doc/Sedgewick-Talk-Penn2008.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/petar/GoLLRB/ae3b015fd3e9d3cdb1e0b0e3da2caf14dd3f908f/doc/Sedgewick-Talk-Penn2008.pdf -------------------------------------------------------------------------------- /example/ex1.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "github.com/petar/GoLLRB/llrb" 6 | ) 7 | 8 | func lessInt(a, b interface{}) bool { return a.(int) < b.(int) } 9 | 10 | func main() { 11 | tree := llrb.New(lessInt) 12 | tree.ReplaceOrInsert(1) 13 | tree.ReplaceOrInsert(2) 14 | tree.ReplaceOrInsert(3) 15 | tree.ReplaceOrInsert(4) 16 | tree.DeleteMin() 17 | tree.Delete(4) 18 | c := tree.IterAscend() 19 | for { 20 | u := <-c 21 | if u == nil { 22 | break 23 | } 24 | fmt.Printf("%d\n", int(u.(int))) 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/petar/GoLLRB 2 | 3 | go 1.15 4 | 5 | require ( 6 | ) 7 | -------------------------------------------------------------------------------- /llrb/avgvar.go: -------------------------------------------------------------------------------- 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 | package llrb 6 | 7 | import "math" 8 | 9 | // avgVar maintains the average and variance of a stream of numbers 10 | // in a space-efficient manner. 11 | type avgVar struct { 12 | count int64 13 | sum, sumsq float64 14 | } 15 | 16 | func (av *avgVar) Init() { 17 | av.count = 0 18 | av.sum = 0.0 19 | av.sumsq = 0.0 20 | } 21 | 22 | func (av *avgVar) Add(sample float64) { 23 | av.count++ 24 | av.sum += sample 25 | av.sumsq += sample * sample 26 | } 27 | 28 | func (av *avgVar) GetCount() int64 { return av.count } 29 | 30 | func (av *avgVar) GetAvg() float64 { return av.sum / float64(av.count) } 31 | 32 | func (av *avgVar) GetTotal() float64 { return av.sum } 33 | 34 | func (av *avgVar) GetVar() float64 { 35 | a := av.GetAvg() 36 | return av.sumsq/float64(av.count) - a*a 37 | } 38 | 39 | func (av *avgVar) GetStdDev() float64 { return math.Sqrt(av.GetVar()) } 40 | -------------------------------------------------------------------------------- /llrb/iterator.go: -------------------------------------------------------------------------------- 1 | package llrb 2 | 3 | type ItemIterator func(i Item) bool 4 | 5 | //func (t *Tree) Ascend(iterator ItemIterator) { 6 | // t.AscendGreaterOrEqual(Inf(-1), iterator) 7 | //} 8 | 9 | func (t *LLRB) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) { 10 | t.ascendRange(t.root, greaterOrEqual, lessThan, iterator) 11 | } 12 | 13 | func (t *LLRB) ascendRange(h *Node, inf, sup Item, iterator ItemIterator) bool { 14 | if h == nil { 15 | return true 16 | } 17 | if !less(h.Item, sup) { 18 | return t.ascendRange(h.Left, inf, sup, iterator) 19 | } 20 | if less(h.Item, inf) { 21 | return t.ascendRange(h.Right, inf, sup, iterator) 22 | } 23 | 24 | if !t.ascendRange(h.Left, inf, sup, iterator) { 25 | return false 26 | } 27 | if !iterator(h.Item) { 28 | return false 29 | } 30 | return t.ascendRange(h.Right, inf, sup, iterator) 31 | } 32 | 33 | // AscendGreaterOrEqual will call iterator once for each element greater or equal to 34 | // pivot in ascending order. It will stop whenever the iterator returns false. 35 | func (t *LLRB) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) { 36 | t.ascendGreaterOrEqual(t.root, pivot, iterator) 37 | } 38 | 39 | func (t *LLRB) ascendGreaterOrEqual(h *Node, pivot Item, iterator ItemIterator) bool { 40 | if h == nil { 41 | return true 42 | } 43 | if !less(h.Item, pivot) { 44 | if !t.ascendGreaterOrEqual(h.Left, pivot, iterator) { 45 | return false 46 | } 47 | if !iterator(h.Item) { 48 | return false 49 | } 50 | } 51 | return t.ascendGreaterOrEqual(h.Right, pivot, iterator) 52 | } 53 | 54 | // AscendLessThan will call iterator once for each element lower than 55 | // pivot in ascending order. It will stop whenever the iterator returns false. 56 | func (t *LLRB) AscendLessThan(pivot Item, iterator ItemIterator) { 57 | t.ascendLessThan(t.root, pivot, iterator) 58 | } 59 | 60 | func (t *LLRB) ascendLessThan(h *Node, pivot Item, iterator ItemIterator) bool { 61 | if h == nil { 62 | return true 63 | } 64 | if !t.ascendLessThan(h.Left, pivot, iterator) { 65 | return false 66 | } 67 | if less(h.Item, pivot) { 68 | if !iterator(h.Item) { 69 | return false 70 | } 71 | return t.ascendLessThan(h.Right, pivot, iterator) 72 | } 73 | return true 74 | } 75 | 76 | // DescendLessOrEqual will call iterator once for each element less than the 77 | // pivot in descending order. It will stop whenever the iterator returns false. 78 | func (t *LLRB) DescendLessOrEqual(pivot Item, iterator ItemIterator) { 79 | t.descendLessOrEqual(t.root, pivot, iterator) 80 | } 81 | 82 | func (t *LLRB) descendLessOrEqual(h *Node, pivot Item, iterator ItemIterator) bool { 83 | if h == nil { 84 | return true 85 | } 86 | if less(h.Item, pivot) || !less(pivot, h.Item) { 87 | if !t.descendLessOrEqual(h.Right, pivot, iterator) { 88 | return false 89 | } 90 | if !iterator(h.Item) { 91 | return false 92 | } 93 | } 94 | return t.descendLessOrEqual(h.Left, pivot, iterator) 95 | } 96 | -------------------------------------------------------------------------------- /llrb/iterator_test.go: -------------------------------------------------------------------------------- 1 | package llrb 2 | 3 | import ( 4 | "reflect" 5 | "testing" 6 | ) 7 | 8 | func TestAscendGreaterOrEqual(t *testing.T) { 9 | tree := New() 10 | tree.InsertNoReplace(Int(4)) 11 | tree.InsertNoReplace(Int(6)) 12 | tree.InsertNoReplace(Int(1)) 13 | tree.InsertNoReplace(Int(3)) 14 | var ary []Item 15 | tree.AscendGreaterOrEqual(Int(-1), func(i Item) bool { 16 | ary = append(ary, i) 17 | return true 18 | }) 19 | expected := []Item{Int(1), Int(3), Int(4), Int(6)} 20 | if !reflect.DeepEqual(ary, expected) { 21 | t.Errorf("expected %v but got %v", expected, ary) 22 | } 23 | ary = nil 24 | tree.AscendGreaterOrEqual(Int(3), func(i Item) bool { 25 | ary = append(ary, i) 26 | return true 27 | }) 28 | expected = []Item{Int(3), Int(4), Int(6)} 29 | if !reflect.DeepEqual(ary, expected) { 30 | t.Errorf("expected %v but got %v", expected, ary) 31 | } 32 | ary = nil 33 | tree.AscendGreaterOrEqual(Int(2), func(i Item) bool { 34 | ary = append(ary, i) 35 | return true 36 | }) 37 | expected = []Item{Int(3), Int(4), Int(6)} 38 | if !reflect.DeepEqual(ary, expected) { 39 | t.Errorf("expected %v but got %v", expected, ary) 40 | } 41 | } 42 | 43 | func TestDescendLessOrEqual(t *testing.T) { 44 | tree := New() 45 | tree.InsertNoReplace(Int(4)) 46 | tree.InsertNoReplace(Int(6)) 47 | tree.InsertNoReplace(Int(1)) 48 | tree.InsertNoReplace(Int(3)) 49 | var ary []Item 50 | tree.DescendLessOrEqual(Int(10), func(i Item) bool { 51 | ary = append(ary, i) 52 | return true 53 | }) 54 | expected := []Item{Int(6), Int(4), Int(3), Int(1)} 55 | if !reflect.DeepEqual(ary, expected) { 56 | t.Errorf("expected %v but got %v", expected, ary) 57 | } 58 | ary = nil 59 | tree.DescendLessOrEqual(Int(4), func(i Item) bool { 60 | ary = append(ary, i) 61 | return true 62 | }) 63 | expected = []Item{Int(4), Int(3), Int(1)} 64 | if !reflect.DeepEqual(ary, expected) { 65 | t.Errorf("expected %v but got %v", expected, ary) 66 | } 67 | ary = nil 68 | tree.DescendLessOrEqual(Int(5), func(i Item) bool { 69 | ary = append(ary, i) 70 | return true 71 | }) 72 | expected = []Item{Int(4), Int(3), Int(1)} 73 | if !reflect.DeepEqual(ary, expected) { 74 | t.Errorf("expected %v but got %v", expected, ary) 75 | } 76 | } 77 | 78 | func TestAscendLessThan(t *testing.T) { 79 | tree := New() 80 | tree.InsertNoReplace(Int(4)) 81 | tree.InsertNoReplace(Int(6)) 82 | tree.InsertNoReplace(Int(1)) 83 | tree.InsertNoReplace(Int(3)) 84 | var ary []Item 85 | tree.AscendLessThan(Int(10), func(i Item) bool { 86 | ary = append(ary, i) 87 | return true 88 | }) 89 | expected := []Item{Int(1), Int(3), Int(4), Int(6)} 90 | if !reflect.DeepEqual(ary, expected) { 91 | t.Errorf("expected %v but got %v", expected, ary) 92 | } 93 | ary = nil 94 | tree.AscendLessThan(Int(4), func(i Item) bool { 95 | ary = append(ary, i) 96 | return true 97 | }) 98 | expected = []Item{Int(1), Int(3)} 99 | if !reflect.DeepEqual(ary, expected) { 100 | t.Errorf("expected %v but got %v", expected, ary) 101 | } 102 | ary = nil 103 | tree.AscendLessThan(Int(5), func(i Item) bool { 104 | ary = append(ary, i) 105 | return true 106 | }) 107 | expected = []Item{Int(1), Int(3), Int(4)} 108 | if !reflect.DeepEqual(ary, expected) { 109 | t.Errorf("expected %v but got %v", expected, ary) 110 | } 111 | } 112 | -------------------------------------------------------------------------------- /llrb/llrb-stats.go: -------------------------------------------------------------------------------- 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 | package llrb 6 | 7 | // GetHeight returns an item in the tree with key @key, and it's height in the tree 8 | func (t *LLRB) GetHeight(key Item) (result Item, depth int) { 9 | return t.getHeight(t.root, key) 10 | } 11 | 12 | func (t *LLRB) getHeight(h *Node, item Item) (Item, int) { 13 | if h == nil { 14 | return nil, 0 15 | } 16 | if less(item, h.Item) { 17 | result, depth := t.getHeight(h.Left, item) 18 | return result, depth + 1 19 | } 20 | if less(h.Item, item) { 21 | result, depth := t.getHeight(h.Right, item) 22 | return result, depth + 1 23 | } 24 | return h.Item, 0 25 | } 26 | 27 | // HeightStats returns the average and standard deviation of the height 28 | // of elements in the tree 29 | func (t *LLRB) HeightStats() (avg, stddev float64) { 30 | av := &avgVar{} 31 | heightStats(t.root, 0, av) 32 | return av.GetAvg(), av.GetStdDev() 33 | } 34 | 35 | func heightStats(h *Node, d int, av *avgVar) { 36 | if h == nil { 37 | return 38 | } 39 | av.Add(float64(d)) 40 | if h.Left != nil { 41 | heightStats(h.Left, d+1, av) 42 | } 43 | if h.Right != nil { 44 | heightStats(h.Right, d+1, av) 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /llrb/llrb.go: -------------------------------------------------------------------------------- 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 | } 457 | -------------------------------------------------------------------------------- /llrb/llrb_test.go: -------------------------------------------------------------------------------- 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 | package llrb 6 | 7 | import ( 8 | "math" 9 | "math/rand" 10 | "testing" 11 | ) 12 | 13 | func TestCases(t *testing.T) { 14 | tree := New() 15 | tree.ReplaceOrInsert(Int(1)) 16 | tree.ReplaceOrInsert(Int(1)) 17 | if tree.Len() != 1 { 18 | t.Errorf("expecting len 1") 19 | } 20 | if !tree.Has(Int(1)) { 21 | t.Errorf("expecting to find key=1") 22 | } 23 | 24 | tree.Delete(Int(1)) 25 | if tree.Len() != 0 { 26 | t.Errorf("expecting len 0") 27 | } 28 | if tree.Has(Int(1)) { 29 | t.Errorf("not expecting to find key=1") 30 | } 31 | 32 | tree.Delete(Int(1)) 33 | if tree.Len() != 0 { 34 | t.Errorf("expecting len 0") 35 | } 36 | if tree.Has(Int(1)) { 37 | t.Errorf("not expecting to find key=1") 38 | } 39 | } 40 | 41 | func TestReverseInsertOrder(t *testing.T) { 42 | tree := New() 43 | n := 100 44 | for i := 0; i < n; i++ { 45 | tree.ReplaceOrInsert(Int(n - i)) 46 | } 47 | i := 0 48 | tree.AscendGreaterOrEqual(Int(0), func(item Item) bool { 49 | i++ 50 | if item.(Int) != Int(i) { 51 | t.Errorf("bad order: got %d, expect %d", item.(Int), i) 52 | } 53 | return true 54 | }) 55 | } 56 | 57 | func TestRange(t *testing.T) { 58 | tree := New() 59 | order := []String{ 60 | "ab", "aba", "abc", "a", "aa", "aaa", "b", "a-", "a!", 61 | } 62 | for _, i := range order { 63 | tree.ReplaceOrInsert(i) 64 | } 65 | k := 0 66 | tree.AscendRange(String("ab"), String("ac"), func(item Item) bool { 67 | if k > 3 { 68 | t.Fatalf("returned more items than expected") 69 | } 70 | i1 := order[k] 71 | i2 := item.(String) 72 | if i1 != i2 { 73 | t.Errorf("expecting %s, got %s", i1, i2) 74 | } 75 | k++ 76 | return true 77 | }) 78 | } 79 | 80 | func TestRandomInsertOrder(t *testing.T) { 81 | tree := New() 82 | n := 1000 83 | perm := rand.Perm(n) 84 | for i := 0; i < n; i++ { 85 | tree.ReplaceOrInsert(Int(perm[i])) 86 | } 87 | j := 0 88 | tree.AscendGreaterOrEqual(Int(0), func(item Item) bool { 89 | if item.(Int) != Int(j) { 90 | t.Fatalf("bad order") 91 | } 92 | j++ 93 | return true 94 | }) 95 | } 96 | 97 | func TestRandomReplace(t *testing.T) { 98 | tree := New() 99 | n := 100 100 | perm := rand.Perm(n) 101 | for i := 0; i < n; i++ { 102 | tree.ReplaceOrInsert(Int(perm[i])) 103 | } 104 | perm = rand.Perm(n) 105 | for i := 0; i < n; i++ { 106 | if replaced := tree.ReplaceOrInsert(Int(perm[i])); replaced == nil || replaced.(Int) != Int(perm[i]) { 107 | t.Errorf("error replacing") 108 | } 109 | } 110 | } 111 | 112 | func TestRandomInsertSequentialDelete(t *testing.T) { 113 | tree := New() 114 | n := 1000 115 | perm := rand.Perm(n) 116 | for i := 0; i < n; i++ { 117 | tree.ReplaceOrInsert(Int(perm[i])) 118 | } 119 | for i := 0; i < n; i++ { 120 | tree.Delete(Int(i)) 121 | } 122 | } 123 | 124 | func TestRandomInsertDeleteNonExistent(t *testing.T) { 125 | tree := New() 126 | n := 100 127 | perm := rand.Perm(n) 128 | for i := 0; i < n; i++ { 129 | tree.ReplaceOrInsert(Int(perm[i])) 130 | } 131 | if tree.Delete(Int(200)) != nil { 132 | t.Errorf("deleted non-existent item") 133 | } 134 | if tree.Delete(Int(-2)) != nil { 135 | t.Errorf("deleted non-existent item") 136 | } 137 | for i := 0; i < n; i++ { 138 | if u := tree.Delete(Int(i)); u == nil || u.(Int) != Int(i) { 139 | t.Errorf("delete failed") 140 | } 141 | } 142 | if tree.Delete(Int(200)) != nil { 143 | t.Errorf("deleted non-existent item") 144 | } 145 | if tree.Delete(Int(-2)) != nil { 146 | t.Errorf("deleted non-existent item") 147 | } 148 | } 149 | 150 | func TestRandomInsertPartialDeleteOrder(t *testing.T) { 151 | tree := New() 152 | n := 100 153 | perm := rand.Perm(n) 154 | for i := 0; i < n; i++ { 155 | tree.ReplaceOrInsert(Int(perm[i])) 156 | } 157 | for i := 1; i < n-1; i++ { 158 | tree.Delete(Int(i)) 159 | } 160 | j := 0 161 | tree.AscendGreaterOrEqual(Int(0), func(item Item) bool { 162 | switch j { 163 | case 0: 164 | if item.(Int) != Int(0) { 165 | t.Errorf("expecting 0") 166 | } 167 | case 1: 168 | if item.(Int) != Int(n-1) { 169 | t.Errorf("expecting %d", n-1) 170 | } 171 | } 172 | j++ 173 | return true 174 | }) 175 | } 176 | 177 | func TestRandomInsertStats(t *testing.T) { 178 | tree := New() 179 | n := 100000 180 | perm := rand.Perm(n) 181 | for i := 0; i < n; i++ { 182 | tree.ReplaceOrInsert(Int(perm[i])) 183 | } 184 | avg, _ := tree.HeightStats() 185 | expAvg := math.Log2(float64(n)) - 1.5 186 | if math.Abs(avg-expAvg) >= 2.0 { 187 | t.Errorf("too much deviation from expected average height") 188 | } 189 | } 190 | 191 | func BenchmarkInsert(b *testing.B) { 192 | tree := New() 193 | for i := 0; i < b.N; i++ { 194 | tree.ReplaceOrInsert(Int(b.N - i)) 195 | } 196 | } 197 | 198 | func BenchmarkDelete(b *testing.B) { 199 | b.StopTimer() 200 | tree := New() 201 | for i := 0; i < b.N; i++ { 202 | tree.ReplaceOrInsert(Int(b.N - i)) 203 | } 204 | b.StartTimer() 205 | for i := 0; i < b.N; i++ { 206 | tree.Delete(Int(i)) 207 | } 208 | } 209 | 210 | func BenchmarkDeleteMin(b *testing.B) { 211 | b.StopTimer() 212 | tree := New() 213 | for i := 0; i < b.N; i++ { 214 | tree.ReplaceOrInsert(Int(b.N - i)) 215 | } 216 | b.StartTimer() 217 | for i := 0; i < b.N; i++ { 218 | tree.DeleteMin() 219 | } 220 | } 221 | 222 | func TestInsertNoReplace(t *testing.T) { 223 | tree := New() 224 | n := 1000 225 | for q := 0; q < 2; q++ { 226 | perm := rand.Perm(n) 227 | for i := 0; i < n; i++ { 228 | tree.InsertNoReplace(Int(perm[i])) 229 | } 230 | } 231 | j := 0 232 | tree.AscendGreaterOrEqual(Int(0), func(item Item) bool { 233 | if item.(Int) != Int(j/2) { 234 | t.Fatalf("bad order") 235 | } 236 | j++ 237 | return true 238 | }) 239 | } 240 | -------------------------------------------------------------------------------- /llrb/util.go: -------------------------------------------------------------------------------- 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 | package llrb 6 | 7 | type Int int 8 | 9 | func (x Int) Less(than Item) bool { 10 | return x < than.(Int) 11 | } 12 | 13 | type String string 14 | 15 | func (x String) Less(than Item) bool { 16 | return x < than.(String) 17 | } 18 | --------------------------------------------------------------------------------