Swift binary search tree

binary search tree.swift

class Node : NSObject {
    
    var index : Int = -1
    
    var data : AnyObject?
    
    var left : Node? = nil
    var right : Node? = nil
    
    init(index : Int, data : AnyObject?) {
        self.index = index
        if let data = data {
            self.data = data
        }
    }
    
    init(index : Int, data : AnyObject? = nil, left : Node , right: Node) {
        self.index = index
        self.left = left
        self.right = right
        if let data = data {
            self.data = data
        }
    }
    
}

class BST : NSObject {
    
    var root : Node!
 
    // Insert Function
    
    func insert(index : Int , withData data : AnyObject? = nil) {
        
        let node = Node(index : index, data: data )
        
        if root == nil {
            root = node
        } else {
            
            if let findNode = find(index) {
                findNode.data = data
                print("Replacing node with index \(index). Node with index already existe")
                return
            }
            
            var currentNode = root
            
            var parent : Node? = nil
            
            var didFinishInserting = false
            while !didFinishInserting {
                parent = currentNode
                if index < currentNode.index {
                    currentNode = currentNode.left
                    if currentNode == nil {
                        parent!.left = node
                        didFinishInserting = true
                    }
                } else {
                    currentNode = currentNode.right
                    if currentNode == nil {
                        parent!.right = node
                        didFinishInserting = true
                    }
                }
            }
        }
    }
    
    // Traversal Functions
    
    func inOrder(node : Node?) {
        if let node = node {
            inOrder(node.left)
            print(node.index)
            inOrder(node.right)
        }
    }
    
    func preOrder(node : Node?) {
        if let node = node {
            print(node.index)
            preOrder(node.left)
            preOrder(node.right)
        }
    }
    
    func postOrder(node : Node?) {
        if let node = node {
            postOrder(node.left)
            postOrder(node.right)
            print(node.index)
        }
    }
    
    // Search Functions
    
    func getMin() -> Node {
        var current = self.root
        while current.left != nil {
            current = current.left
        }
        return current
    }
    
    func getMax() -> Node {
        var current = self.root
        while current.right != nil {
            current = current.right
        }
        return current
    }
    
    func find(index :Int) -> Node? {
        var current = self.root
        while current != nil && current.index != index {
            if index < current.index {
                current = current.left
            } else {
                current = current.right
            }
        }
        return current
    }
    
    // Remove Functions
    
    func removeIndex (index :Int) {
        
        func getSmallest(node : Node) -> Node {
            if let leftNode = node.left  {
                return getSmallest(leftNode)
            } else {
                return node
            }
        }
        
        func removeNode(node : Node?, index : Int) -> Node?  {
            if let node = node {
                if index == node.index {
                    
                    if node.left == nil && node.right == nil {
                        return nil
                    }
                    
                    if node.left == nil {
                        return node.right
                    }
                    
                    if node.right == nil {
                        return node.left
                    }
                    
                    let tmpNode = getSmallest(node.right!)
                    node.index = tmpNode.index
                    node.right = removeNode(node.right, index: tmpNode.index)
                    
                    return node
                } else if index < node.index {
                    node.left = removeNode(node.left, index: index)
                    return node
                } else {
                    node.right = removeNode(node.right, index: index)
                    return node
                }
                
            } else {
                return nil
            }
            
        }
        root = removeNode(root, index: index)
    }
}