Iterative Implementation of Binary Search Tree - Depth First Search (DFS) in Swift 4 - In order, pre order, and post order.

BinaryTree-DFS-Iterative-Swift4.swift

class TreeNode: Hashable {
    var val: Int
    var left: TreeNode?
    var right: TreeNode?
    
    init(_ val: Int, _ left: TreeNode? = nil, _ right: TreeNode? = nil) {
        self.val = val
        self.left = left
        self.right = right
    }
    
    static func == (lhs: TreeNode, rhs: TreeNode) -> Bool {
        return lhs === rhs
    }
    
    var hashValue: Int {
        return val.hashValue
    }
}


// Post order traversal
func postOrderDFSIterative(_ root: TreeNode?) {
    guard let root = root else { return }
    var seen: Set<TreeNode> = [root]
    var stack = [root]
    
    while !stack.isEmpty {
        let top = stack.last!
        if let left = top.left, !seen.contains(left) {
            seen.insert(left)
            stack.append(left)
        } else if let right = top.right, !seen.contains(right) {
            seen.insert(right)
            stack.append(right)
        } else {
            print(top.val, terminator: " ")
            stack.removeLast()
        }
    }
    print("")
}

// In order traversal
func inOrderDFSIterative(_ root: TreeNode?) {
    guard let root = root else { return }
    var stack = [TreeNode]()
    var current: TreeNode? = root
    
    while current != nil || !stack.isEmpty {
        if let cur = current {
            stack.append(cur)
            current = cur.left
        } else {
            print(stack.last!.val, terminator: " ")
            current = stack.removeLast().right
        }
    }
    
    print("")
}

// Pre order traversal
func preOrderDFSIterative(_ root: TreeNode?) {
    guard let root = root else { return }
    var stack = [TreeNode]()
    var current: TreeNode? = root
    
    while current != nil || !stack.isEmpty {
        if let cur = current {
            print(cur.val, terminator: " ")
            stack.append(cur)
            current = cur.left
        } else {
            current = stack.removeLast().right
        }
    }
    
    print("")
}