Merge pull request #84 from prankymat/improve-avl-tree-test

added more test cases to test the validity of the AVL Tree

Author: hollance

AVL Tree/AVLTree.playground/Sources/AVLTree.swift

@@ -26,8 +26,8 @@ public class TreeNode<Key: Comparable, Payload> {
   public var payload: Payload?
 
   private var key: Key
-  private var leftChild: Node?
-  private var rightChild: Node?
+  internal var leftChild: Node?
+  internal var rightChild: Node?
   private var height: Int
   weak private var parent: Node?
 
@@ -343,6 +343,8 @@ extension AVLTree {
         } else if node.isRightChild {
           parent.rightChild = nil
         }
+
+        balance(parent)
       } else {
         // at root
         root = nil

AVL Tree/AVLTree.swift

@@ -26,8 +26,8 @@ public class TreeNode<Key: Comparable, Payload> {
   public var payload: Payload?
 
   private var key: Key
-  private var leftChild: Node?
-  private var rightChild: Node?
+  internal var leftChild: Node?
+  internal var rightChild: Node?
   private var height: Int
   weak private var parent: Node?
 
@@ -343,6 +343,8 @@ extension AVLTree {
         } else if node.isRightChild {
           parent.rightChild = nil
         }
+
+        balance(parent)
       } else {
         // at root
         root = nil

AVL Tree/Tests/AVLTreeTests.swift

@@ -11,39 +11,61 @@ import XCTest
 class AVLTreeTests: XCTestCase {
 
   var tree : AVLTree<Int, String>?
-  
+
   override func setUp() {
     super.setUp()
 
     tree = AVLTree()
-    tree?.insert(1, "A")
-    tree?.insert(2, "B")
-    tree?.insert(3, "C")
-    tree?.insert(4, "D")
   }
 
   override func tearDown() {
     // Put teardown code here. This method is called after the invocation of each test method in the class.
     super.tearDown()
   }
-  
+
+  func testAVLTreeBalancedAutoPopulate() {
+    self.tree?.autopopulateWithNodes(10)
+
+    do {
+      try self.tree?.inOrderCheckBalanced(self.tree?.root)
+    } catch _ {
+      XCTFail("Tree is not balanced after autopopulate")
+    }
+  }
+
+  func testAVLTreeBalancedInsert() {
+    self.tree?.autopopulateWithNodes(5)
+
+    for i in 6...10 {
+      self.tree?.insert(i)
+      do {
+        try self.tree?.inOrderCheckBalanced(self.tree?.root)
+      } catch _ {
+        XCTFail("Tree is not balanced after inserting " + String(i))
+      }
+    }
+  }
+
+  func testAVLTreeBalancedDelete() {
+    self.tree?.autopopulateWithNodes(5)
+
+    for i in 1...6 {
+      self.tree?.delete(i)
+      do {
+        try self.tree?.inOrderCheckBalanced(self.tree?.root)
+      } catch _ {
+        XCTFail("Tree is not balanced after deleting " + String(i))
+      }
+    }
+  }
+
   func testEmptyInitialization() {
     let tree = AVLTree<Int,String>()
 
     XCTAssertEqual(tree.size, 0)
     XCTAssertNil(tree.root)
   }
-  
-  func testNotEmptyInitialization() {
-    XCTAssertNotNil(self.tree?.root)
-    XCTAssertNotEqual(self.tree!.size, 0)
-  }
-  
-  func testInsertDuplicated() {
-    self.tree?.insert(1, "A")
-    XCTAssertEqual(self.tree?.size, 5, "Duplicated elements should be allowed")
-  }
-  
+
   func testSingleInsertionPerformance() {
     self.measureBlock {
       self.tree?.insert(5, "E")
@@ -92,6 +114,14 @@ class AVLTreeTests: XCTestCase {
     self.tree?.delete(1056)
     XCTAssertNil(self.tree?.search(1056), "Key should not exist")
   }
+
+  func testInsertSize() {
+    let tree = AVLTree<Int, String>()
+    for i in 0...5 {
+      tree.insert(i, "")
+      XCTAssertEqual(tree.size, i + 1, "Insert didn't update size correctly!")
+    }
+  }
 
   func testDelete() {
     let permutations = [
@@ -109,9 +139,12 @@ class AVLTreeTests: XCTestCase {
       tree.insert(3, "three")
       tree.insert(4, "two")
       tree.insert(5, "one")
-      
+
+      var count = tree.size
       for i in p {
         tree.delete(i)
+        count -= 1
+        XCTAssertEqual(tree.size, count, "Delete didn't update size correctly!")
       }
     }
   }
@@ -126,3 +159,31 @@ extension AVLTree where Key : SignedIntegerType {
     }
   }
 }
+
+enum AVLTreeError: ErrorType {
+  case NotBalanced
+}
+
+extension AVLTree where Key : SignedIntegerType {
+  func height(node: Node?) -> Int {
+    if let node = node {
+      let lHeight = height(node.leftChild)
+      let rHeight = height(node.rightChild)
+
+      return max(lHeight, rHeight) + 1
+    }
+    return 0
+  }
+
+  func inOrderCheckBalanced(node: Node?) throws {
+    if let node = node {
+      guard abs(height(node.leftChild) - height(node.rightChild)) <= 1 else {
+        throw AVLTreeError.NotBalanced
+      }
+      try inOrderCheckBalanced(node.leftChild)
+      try inOrderCheckBalanced(node.rightChild)
+    }
+  }
+}
+
+