001/* 002 * Licensed to the Apache Software Foundation (ASF) under one 003 * or more contributor license agreements. See the NOTICE file 004 * distributed with this work for additional information 005 * regarding copyright ownership. The ASF licenses this file 006 * to you under the Apache License, Version 2.0 (the 007 * "License"); you may not use this file except in compliance 008 * with the License. You may obtain a copy of the License at 009 * 010 * http://www.apache.org/licenses/LICENSE-2.0 011 * 012 * Unless required by applicable law or agreed to in writing, 013 * software distributed under the License is distributed on an 014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 015 * KIND, either express or implied. See the License for the 016 * specific language governing permissions and limitations 017 * under the License. 018 * 019 */ 020package org.apache.directory.server.core.avltree.avl; 021 022 023import java.util.Iterator; 024import java.util.Stack; 025 026 027/** 028 * AVL Tree Set 029 * 030 * @author Vladimir Lysyy (http://bobah.net) 031 * 032 */ 033public class AvlTreeSet<T extends Comparable<T>> implements Iterable<T> 034{ 035 private AvlNode<T> tree; 036 private int size = 0; 037 038 private final boolean useFreeList; 039 040 private Stack<AvlNode<T>> freeList = new Stack<>(); 041 042 043 public AvlTreeSet() 044 { 045 this( false ); 046 } 047 048 049 public AvlTreeSet( boolean useFreeList ) 050 { 051 this.useFreeList = useFreeList; 052 } 053 054 055 public final int height() 056 { 057 return ( tree == null ) ? 0 : tree.height + 1; 058 } 059 060 061 public final int size() 062 { 063 return size; 064 } 065 066 067 public final Iterator<T> iterator() 068 { 069 return new AvlTreeIterator<>( tree ); 070 } 071 072 073 public final boolean insert( T value ) 074 { 075 // empty tree case 076 if ( tree == null ) 077 { 078 tree = newNode( null, value ); 079 ++size; 080 081 return true; 082 } 083 084 AvlNode<T> node = tree; 085 086 // find the place and insert the value 087 int cmp = value.compareTo( node.value ); 088 089 for ( ; cmp != 0; cmp = value.compareTo( node.value ) ) 090 { 091 if ( cmp < 0 ) 092 { 093 if ( node.left == null ) 094 { 095 node.left = newNode( node, value ); 096 break; 097 } 098 else 099 { 100 node = node.left; 101 } 102 } 103 else if ( cmp > 0 ) 104 { 105 if ( node.right == null ) 106 { 107 node.right = newNode( node, value ); 108 break; 109 } 110 else 111 { 112 node = node.right; 113 } 114 } 115 else 116 { 117 assert false : "should never happen"; 118 } 119 } 120 121 // node with _value_ already exists 122 if ( cmp == 0 ) 123 { 124 return false; 125 } 126 127 rebalanceUp( node ); 128 ++size; 129 130 return true; 131 } 132 133 134 private AvlNode<T> newNode( AvlNode<T> parent, T value ) 135 { 136 if ( !useFreeList || freeList.isEmpty() ) 137 { 138 return new AvlNode<>( parent, value ); 139 } 140 else 141 { 142 AvlNode<T> node = freeList.pop(); 143 144 return node.reset( parent, value ); 145 } 146 } 147 148 149 private void recycleNode( AvlNode<T> node ) 150 { 151 if ( !useFreeList ) 152 { 153 return; 154 } 155 156 // keep free list size not bigger than tree size 157 while ( freeList.size() > size ) 158 { 159 freeList.pop(); 160 } 161 162 if ( freeList.size() == size ) 163 { 164 return; 165 } 166 167 freeList.push( node ); 168 } 169 170 171 private void rebalanceUp( AvlNode<T> node ) 172 { 173 while ( node != null ) 174 { 175 int heightBefore = node.height; 176 updateHeight( node ); 177 178 // rebalance 179 if ( node.balance == -2 ) 180 { 181 node = bigRightRotation( node ); 182 } 183 else if ( node.balance == 2 ) 184 { 185 node = bigLeftRotation( node ); 186 } 187 188 if ( node.parent == null ) 189 { 190 tree = node; 191 } 192 193 // if parent node is not affected 194 if ( heightBefore == node.height ) 195 { 196 break; 197 } 198 199 node = node.parent; 200 } 201 } 202 203 204 public final boolean remove( T value ) 205 { 206 AvlNode<T> node = tree; 207 208 if ( node == null ) 209 { 210 return false; 211 } 212 213 // find the node to be removed 214 for ( int cmp = value.compareTo( node.value ); cmp != 0; cmp = value.compareTo( node.value ) ) 215 { 216 node = ( cmp < 0 ) ? node.left : node.right; 217 218 if ( node == null ) 219 { 220 return false; 221 } 222 } 223 224 // find a replacement node (if needed) 225 final int left = -1; 226 final int right = 1; 227 final int none = 0; 228 int replaceFrom = none; 229 230 if ( node.left != null && node.right == null ) 231 { 232 replaceFrom = left; 233 } 234 else if ( node.right != null && node.left == null ) 235 { 236 replaceFrom = right; 237 } 238 else if ( node.right != null && node.left != null ) 239 { 240 if ( node.balance < 0 ) 241 { 242 replaceFrom = left; 243 } 244 else if ( node.balance > 0 ) 245 { 246 replaceFrom = right; 247 } 248 else 249 { 250 replaceFrom = left; // TODO: asymmetry 251 } 252 } 253 else 254 { // node is itself a leaf, replacement is not needed 255 if ( node.parent == null ) 256 { // the tree root, single node in the tree 257 tree = null; 258 --size; 259 recycleNode( node ); 260 261 return true; 262 } 263 else 264 { // non-root leaf node 265 // detach from parent 266 if ( node.parent.left == node ) 267 { 268 node.parent.left = null; 269 } 270 else 271 { 272 node.parent.right = null; 273 } 274 275 AvlNode<T> dead = node; 276 // update heights/rebalance from node's parents up (the bottom of this method) 277 node = node.parent; 278 recycleNode( dead ); 279 replaceFrom = none; 280 } 281 } 282 283 if ( replaceFrom != none ) 284 { 285 AvlNode<T> leaf = null; 286 287 if ( replaceFrom == left ) 288 { 289 leaf = node.left; 290 291 while ( leaf.left != null || leaf.right != null ) 292 { 293 if ( leaf.right != null ) 294 { 295 leaf = leaf.right; 296 } 297 else 298 { 299 // the rotation should ensure (leaf.right != null) on the next iteration 300 leaf = smallRightRotation( leaf ); 301 } 302 } 303 } 304 else if ( replaceFrom == right ) 305 { 306 leaf = node.right; 307 308 while ( leaf.right != null || leaf.left != null ) 309 { 310 if ( leaf.left != null ) 311 { 312 leaf = leaf.left; 313 } 314 else 315 { 316 // the rotation should ensure (leaf.left != null) on the next iteration 317 leaf = smallLeftRotation( leaf ); 318 } 319 } 320 } 321 else 322 { 323 assert false : "should never happen"; 324 } 325 326 assert leaf != null : "replacement leaf should always exist at this point"; 327 328 // detach leaf from its parent 329 if ( leaf.parent.left == leaf ) 330 { 331 leaf.parent.left = null; 332 } 333 else if ( leaf.parent.right == leaf ) 334 { 335 leaf.parent.right = null; 336 } 337 else 338 { 339 assert false : "broken parent/child reference in the tree"; 340 } 341 342 node.value = leaf.value; // replace node value with leaf's value 343 node = leaf.parent; // change recursion point down so that below down-up update picks up 344 // everything from leaf's parent up 345 346 recycleNode( leaf ); 347 } 348 349 rebalanceUp( node ); 350 351 --size; 352 353 return true; 354 } 355 356 357 public final boolean contains( T value ) 358 { 359 AvlNode<T> node = tree; 360 361 while ( node != null ) 362 { 363 int cmp = value.compareTo( node.value ); 364 365 if ( cmp < 0 ) 366 { 367 node = node.left; 368 } 369 else if ( cmp > 0 ) 370 { 371 node = node.right; 372 } 373 else 374 { 375 return true; 376 } 377 } 378 379 return false; 380 381 } 382 383 384 private static <T extends Comparable<T>> void updateHeight( AvlNode<T> node ) 385 { 386 int leftHeight = ( node.left == null ) ? -1 : node.left.height; 387 int rightHeight = ( node.right == null ) ? -1 : node.right.height; 388 node.height = 1 + ( rightHeight > leftHeight ? rightHeight : leftHeight ); 389 node.balance = rightHeight - leftHeight; 390 } 391 392 393 private static <T extends Comparable<T>> AvlNode<T> smallLeftRotation( AvlNode<T> node ) 394 { 395 assert node.balance > 0 : "null right child in smallLeft"; 396 397 // update child references 398 AvlNode<T> right = node.right; 399 node.right = right.left; 400 right.left = node; 401 402 // update parent references 403 if ( node.right != null ) 404 { 405 node.right.parent = node; 406 } 407 408 right.parent = node.parent; 409 410 if ( right.parent != null ) 411 { 412 if ( right.parent.left == node ) 413 { 414 node.parent.left = right; 415 } 416 else 417 { 418 node.parent.right = right; 419 } 420 } 421 422 node.parent = right; 423 424 updateHeight( node ); 425 updateHeight( right ); 426 427 return right; 428 } 429 430 431 private static <T extends Comparable<T>> AvlNode<T> smallRightRotation( AvlNode<T> node ) 432 { 433 assert node.balance < 0 : "null left child in smallRight"; 434 435 // update child references 436 AvlNode<T> left = node.left; 437 node.left = left.right; 438 left.right = node; 439 440 // update parent references 441 if ( node.left != null ) 442 { 443 node.left.parent = node; 444 } 445 446 left.parent = node.parent; 447 448 if ( left.parent != null ) 449 { 450 if ( left.parent.left == node ) 451 { 452 node.parent.left = left; 453 } 454 else 455 { 456 node.parent.right = left; 457 } 458 } 459 460 node.parent = left; 461 462 updateHeight( node ); 463 updateHeight( left ); 464 465 return left; 466 } 467 468 469 private static <T extends Comparable<T>> AvlNode<T> bigLeftRotation( AvlNode<T> node ) 470 { 471 assert node.right != null : "null right child in bigLeft"; 472 473 if ( node.right.balance < 0 ) 474 { 475 node.right = smallRightRotation( node.right ); 476 } 477 478 updateHeight( node ); 479 480 return smallLeftRotation( node ); 481 } 482 483 484 private static <T extends Comparable<T>> AvlNode<T> bigRightRotation( AvlNode<T> node ) 485 { 486 assert node.left != null : "null right child in bigRight"; 487 488 if ( node.left.balance > 0 ) 489 { 490 node.left = smallLeftRotation( node.left ); 491 } 492 493 updateHeight( node ); 494 495 return smallRightRotation( node ); 496 } 497}