Merge pull request kodecocodes#628 from JulioBBL/master

Aditions to AVLTree

Author: kelvinlauKL

AVL Tree/AVLTree.swift

@@ -22,63 +22,63 @@
 
 public class TreeNode<Key: Comparable, Payload> {
   public typealias Node = TreeNode<Key, Payload>
-
-  var payload: Payload?
-
-  fileprivate var key: Key
+  
+  var payload: Payload?   // Value held by the node
+  
+  fileprivate var key: Key    // Node's name
   internal var leftChild: Node?
   internal var rightChild: Node?
   fileprivate var height: Int
-  fileprivate weak var parent: Node?
-
+  weak fileprivate var parent: Node?
+  
   public init(key: Key, payload: Payload?, leftChild: Node?, rightChild: Node?, parent: Node?, height: Int) {
     self.key = key
     self.payload = payload
     self.leftChild = leftChild
     self.rightChild = rightChild
     self.parent = parent
     self.height = height
-
+    
     self.leftChild?.parent = self
     self.rightChild?.parent = self
   }
-
+  
   public convenience init(key: Key, payload: Payload?) {
     self.init(key: key, payload: payload, leftChild: nil, rightChild: nil, parent: nil, height: 1)
   }
-
+  
   public convenience init(key: Key) {
     self.init(key: key, payload: nil)
   }
-
+  
   var isRoot: Bool {
     return parent == nil
   }
-
+  
   var isLeaf: Bool {
     return rightChild == nil && leftChild == nil
   }
-
+  
   var isLeftChild: Bool {
     return parent?.leftChild === self
   }
-
+  
   var isRightChild: Bool {
     return parent?.rightChild === self
   }
-
+  
   var hasLeftChild: Bool {
     return leftChild != nil
   }
-
+  
   var hasRightChild: Bool {
     return rightChild != nil
   }
-
+  
   var hasAnyChild: Bool {
     return leftChild != nil || rightChild != nil
   }
-
+  
   var hasBothChildren: Bool {
     return leftChild != nil && rightChild != nil
   }
@@ -88,10 +88,10 @@ public class TreeNode<Key: Comparable, Payload> {
 
 open class AVLTree<Key: Comparable, Payload> {
   public typealias Node = TreeNode<Key, Payload>
-
+  
   fileprivate(set) var root: Node?
   fileprivate(set) var size = 0
-
+  
   public init() { }
 }
 
@@ -101,7 +101,7 @@ extension TreeNode {
   public func minimum() -> TreeNode? {
     return leftChild?.minimum() ?? self
   }
-
+  
   public func maximum() -> TreeNode? {
     return rightChild?.maximum() ?? self
   }
@@ -112,11 +112,11 @@ extension AVLTree {
     get { return search(input: key) }
     set { insert(key: key, payload: newValue) }
   }
-
+  
   public func search(input: Key) -> Payload? {
     return search(key: input, node: root)?.payload
   }
-
+  
   fileprivate func search(key: Key, node: Node?) -> Node? {
     if let node = node {
       if key == node.key {
@@ -142,7 +142,7 @@ extension AVLTree {
     }
     size += 1
   }
-
+  
   private func insert(input: Key, payload: Payload?, node: Node) {
     if input < node.key {
       if let child = node.leftChild {
@@ -152,7 +152,7 @@ extension AVLTree {
         node.leftChild = child
         balance(node: child)
       }
-    } else {
+    } else if input != node.key {
       if let child = node.rightChild {
         insert(input: input, payload: payload, node: child)
       } else {
@@ -175,25 +175,25 @@ extension AVLTree {
       updateHeightUpwards(node: node.parent)
     }
   }
-
+  
   fileprivate func lrDifference(node: Node?) -> Int {
     let lHeight = node?.leftChild?.height ?? 0
     let rHeight = node?.rightChild?.height ?? 0
     return lHeight - rHeight
   }
-
+  
   fileprivate func balance(node: Node?) {
     guard let node = node else {
       return
     }
-
+    
     updateHeightUpwards(node: node.leftChild)
     updateHeightUpwards(node: node.rightChild)
-
+    
     var nodes = [Node?](repeating: nil, count: 3)
     var subtrees = [Node?](repeating: nil, count: 4)
     let nodeParent = node.parent
-
+    
     let lrFactor = lrDifference(node: node)
     if lrFactor > 1 {
       // left-left or left-right
@@ -202,7 +202,7 @@ extension AVLTree {
         nodes[0] = node
         nodes[2] = node.leftChild
         nodes[1] = nodes[2]?.leftChild
-
+        
         subtrees[0] = nodes[1]?.leftChild
         subtrees[1] = nodes[1]?.rightChild
         subtrees[2] = nodes[2]?.rightChild
@@ -212,7 +212,7 @@ extension AVLTree {
         nodes[0] = node
         nodes[1] = node.leftChild
         nodes[2] = nodes[1]?.rightChild
-
+        
         subtrees[0] = nodes[1]?.leftChild
         subtrees[1] = nodes[2]?.leftChild
         subtrees[2] = nodes[2]?.rightChild
@@ -225,7 +225,7 @@ extension AVLTree {
         nodes[1] = node
         nodes[2] = node.rightChild
         nodes[0] = nodes[2]?.rightChild
-
+        
         subtrees[0] = nodes[1]?.leftChild
         subtrees[1] = nodes[2]?.leftChild
         subtrees[2] = nodes[0]?.leftChild
@@ -235,7 +235,7 @@ extension AVLTree {
         nodes[1] = node
         nodes[0] = node.rightChild
         nodes[2] = nodes[0]?.leftChild
-
+        
         subtrees[0] = nodes[1]?.leftChild
         subtrees[1] = nodes[2]?.leftChild
         subtrees[2] = nodes[2]?.rightChild
@@ -246,9 +246,9 @@ extension AVLTree {
       balance(node: node.parent)
       return
     }
-
+    
     // nodes[2] is always the head
-
+    
     if node.isRoot {
       root = nodes[2]
       root?.parent = nil
@@ -259,25 +259,25 @@ extension AVLTree {
       nodeParent?.rightChild = nodes[2]
       nodes[2]?.parent = nodeParent
     }
-
+    
     nodes[2]?.leftChild = nodes[1]
     nodes[1]?.parent = nodes[2]
     nodes[2]?.rightChild = nodes[0]
     nodes[0]?.parent = nodes[2]
-
+    
     nodes[1]?.leftChild = subtrees[0]
     subtrees[0]?.parent = nodes[1]
     nodes[1]?.rightChild = subtrees[1]
     subtrees[1]?.parent = nodes[1]
-
+    
     nodes[0]?.leftChild = subtrees[2]
     subtrees[2]?.parent = nodes[0]
     nodes[0]?.rightChild = subtrees[3]
     subtrees[3]?.parent = nodes[0]
-
+    
     updateHeightUpwards(node: nodes[1])    // Update height from left
     updateHeightUpwards(node: nodes[0])    // Update height from right
-
+    
     balance(node: nodes[2]?.parent)
   }
 }
@@ -299,11 +299,35 @@ extension AVLTree {
       display(node: node.leftChild, level: level + 1)
     }
   }
-
+  
   public func display(node: Node) {
     display(node: node, level: 0)
     print("")
   }
+  
+  public func inorder(node: Node?) -> String {
+    var output = ""
+    if let node = node {
+      output = "\(inorder(node: node.leftChild)) \(print("\(node.key) ")) \(inorder(node: node.rightChild))"
+    }
+    return output
+  }
+  
+  public func preorder(node: Node?) -> String {
+    var output = ""
+    if let node = node {
+      output = "\(preorder(node: node.leftChild)) \(print("\(node.key) ")) \(preorder(node: node.rightChild))"
+    }
+    return output
+  }
+  
+  public func postorder(node: Node?) -> String {
+    var output = ""
+    if let node = node {
+      output = "\(postorder(node: node.leftChild)) \(print("\(node.key) ")) \(postorder(node: node.rightChild))"
+    }
+    return output
+  }
 }
 
 // MARK: - Delete node
@@ -318,7 +342,7 @@ extension AVLTree {
       size -= 1
     }
   }
-
+  
   private func delete(node: Node) {
     if node.isLeaf {
       // Just remove and balance up
@@ -327,13 +351,13 @@ extension AVLTree {
           // just in case
           fatalError("Error: tree is invalid.")
         }
-
+        
         if node.isLeftChild {
           parent.leftChild = nil
         } else if node.isRightChild {
           parent.rightChild = nil
         }
-
+        
         balance(node: parent)
       } else {
         // at root
@@ -354,6 +378,46 @@ extension AVLTree {
   }
 }
 
+// MARK: - Advanced Stuff
+
+extension AVLTree {
+  public func doInOrder(node: Node?, _ completion: (Node) -> Void) {
+    if let node = node {
+      doInOrder(node: node.leftChild) { lnode in
+        completion(lnode)
+      }
+      completion(node)
+      doInOrder(node: node.rightChild) { rnode in
+        completion(rnode)
+      }
+    }
+  }
+  
+  public func doInPreOrder(node: Node?, _ completion: (Node) -> Void) {
+    if let node = node {
+      completion(node)
+      doInPreOrder(node: node.leftChild) { lnode in
+        completion(lnode)
+      }
+      doInPreOrder(node: node.rightChild) { rnode in
+        completion(rnode)
+      }
+    }
+  }
+  
+  public func doInPostOrder(node: Node?, _ completion: (Node) -> Void) {
+    if let node = node {
+      doInPostOrder(node: node.leftChild) { lnode in
+        completion(lnode)
+      }
+      doInPostOrder(node: node.rightChild) { rnode in
+        completion(rnode)
+      }
+      completion(node)
+    }
+  }
+}
+
 // MARK: - Debugging
 
 extension TreeNode: CustomDebugStringConvertible {