Binary Search Tree in Swift

binary_search_tree.swift

public indirect enum BinarySearchTree<E: Comparable> {
    public typealias Element = E
    
    case Empty
    case Node(
        left: BinarySearchTree<Element>,
        element: Element,
        right: BinarySearchTree<Element>)
    
    public init() { self = .Empty }
    
    public init(_ element: Element,
        left: BinarySearchTree<Element> = .Empty,
        right: BinarySearchTree<Element> = .Empty)
    {
        self = .Node(left: left, element: element, right: right)
    }
    
    public mutating func insert(element: Element) -> BinarySearchTree<Element> {
        switch self {
        case .Empty:
            self = .Node(left: .Empty, element: element, right: .Empty)
            return self
        case .Node(var left, let nodeElement, var right):
            if element < nodeElement {
                let endNode = left.insert(element)
                self = .Node(left: left, element: nodeElement, right: right)
                return endNode
            } else if element > nodeElement {
                let endNode = right.insert(element)
                self = .Node(left: left, element: nodeElement, right: right)
                return endNode
            } else {
                return self
            }
        }
    }
    
    private mutating func _insertNode(node: BinarySearchTree<Element>)
        -> BinarySearchTree<Element>
    {
        switch (self, node) {
        case (.Empty, .Empty):
            self = node
            return self
        case (.Empty, .Node):
            self = node
            return self
        case (.Node, .Empty):
            return .Empty
        case (.Node(var left, let nodeElement, var right),
            .Node(_ , let insertedElement, _)):
            if insertedElement < nodeElement {
                let endNode = left._insertNode(node)
                self = .Node(left: left, element: nodeElement, right: right)
                return endNode
            } else if insertedElement > nodeElement {
                let endNode = right._insertNode(node)
                self = .Node(left: left, element: nodeElement, right: right)
                return endNode
            } else {
                return self
            }
        }
    }
    
    public mutating func remove(element: Element) -> Element? {
        switch self {
        case .Empty:
            return nil
        case .Node(var left, let nodeElement, var right):
            if element < nodeElement {
                let removed = left.remove(element)
                self = .Node(left: left, element: nodeElement, right: right)
                return removed
            } else if element > nodeElement {
                let removed = right.remove(element)
                self = .Node(left: left, element: nodeElement, right: right)
                return removed
            } else {
                switch (left, right) {
                case (.Empty, Empty):
                    self = .Empty
                    return nodeElement
                case (.Node, .Empty):
                    self = left
                    return nodeElement
                case (.Empty, .Node):
                    self = right
                    return nodeElement
                case (.Node(_,let leftElement, _),
                    .Node(var leftChild, let rightElement, var rightChild)):
                    
                    self = right
                    
                    if leftElement < rightElement {
                        leftChild._insertNode(left)
                        self = .Node(left: leftChild,
                            element: rightElement,
                            right: rightChild)
                    } else if leftElement > rightElement {
                        rightChild._insertNode(left)
                        self = .Node(left: leftChild, 
                            element: rightElement,
                            right: rightChild)
                    } else {
                        
                    }
                    
                    return nodeElement
                }
                
            }
        }
    }
    
    public func contains(element: Element) -> Bool {
        switch self {
        case .Empty:
            return false
        case let .Node(left, nodeElement, right):
            if element < nodeElement {
                return left.contains(element)
            } else if element > nodeElement {
                return right.contains(element)
            } else {
                return true
            }
        }
    }
}

extension BinarySearchTree: SequenceType {
    public func generate() -> AnyGenerator<Element> {
        var stack: [BinarySearchTree] = []
        var current: BinarySearchTree = self
        return anyGenerator { _ -> Element? in
            while true {
                switch current {
                case .Empty where !stack.isEmpty:
                    switch stack.removeLast() {
                    case let .Node(_,element,right):
                        current = right
                        return element
                    default:
                        return nil
                    }
                case .Empty:
                    return nil
                case let .Node(left,_,_):
                    stack.append(current)
                    current = left
                }
            }
        }
    }
}

extension BinarySearchTree: ArrayLiteralConvertible {
    public init(arrayLiteral elements: Element...) {
        self = elements.reduce(BinarySearchTree<Element>.Empty, combine: {
            (var reduced, element) -> BinarySearchTree<Element> in
            reduced.insert(element)
            return reduced
        })
    }
}

var aBST: BinarySearchTree<Int> = [1,6,4,3,7,9,2,0]

for each in aBST {
    each
}

aBST.remove(4)

for each in aBST {
    each
}