balanced BST

balanced_bst.swift

//: Playground - noun: a place where people can play

import UIKit

enum BinarySearchTree<T: Comparable> {
    case empty
    case leaf(T)
    indirect case node(BinarySearchTree<T>, T, BinarySearchTree<T>)

    var height: Int {
        switch self {
        case .empty:
            return 0
        case .leaf:
            return 1
        case let .node(left, _, right):
            return 1 + max(left.height, right.height)
        }
    }
}

extension BinarySearchTree: CustomStringConvertible {
    var description: String {
        switch self {
        case .empty:
            return "."
        case .leaf(let value):
            return "\(value)"
        case let .node(left, value, right):
            return "(\(left.description) <- \(value) -> \(right.description))"
        }
    }
}

extension BinarySearchTree { // MARK: - Traverse
    func traverseInOrder(callout: (T) -> Void) {
        if case let .node(left, value, right) = self {
            left.traverseInOrder(callout: callout)
            callout(value)
            right.traverseInOrder(callout: callout)
        }
    }

    func traverseInOrderWithNodes(callout: (BinarySearchTree<T>) -> Void) {
        if case let .node(left, _, right) = self {
            left.traverseInOrderWithNodes(callout: callout)
            callout(self)
            right.traverseInOrderWithNodes(callout: callout)
        }
    }

    func traversePreOrder(callout: (T) -> Void) {
        if case let .node(left, value, right) = self {
            callout(value)
            left.traverseInOrder(callout: callout)
            right.traverseInOrder(callout: callout)
        }
    }

    func traversePostOrder(callout: (T) -> Void) {
        if case let .node(left, value, right) = self {
            left.traverseInOrder(callout: callout)
            right.traverseInOrder(callout: callout)
            callout(value)
        }
    }
}

extension BinarySearchTree {
    static func sortedArrayToBST(_ sortedArray: [T], start: Int, end: Int) -> BinarySearchTree<T> {
        if start > end {
            return BinarySearchTree.empty
        }

        let middle = (start + end) / 2

        let left = sortedArrayToBST(sortedArray, start: start, end: middle - 1)
        let right = sortedArrayToBST(sortedArray, start: middle + 1, end: end)

        let root = BinarySearchTree.node(left, sortedArray[middle], right)

        return root
    }
}

extension BinarySearchTree {
    func search(x: T) -> BinarySearchTree? {
        switch self {
        case .empty:
            return nil
        case .leaf(let y):
            return x == y ? self : nil
        case let .node(left, y, right):
            if x < y {
                return left.search(x: x)
            } else if x > y {
                return right.search(x: x)
            } else {
                return self
            }
        }
    }
}

let node1 = BinarySearchTree.node(.empty, 1, .empty)
let node2 = BinarySearchTree.node(node1, 2, .empty)
let node3 = BinarySearchTree.node(node2, 3, .empty)
let node7 = BinarySearchTree.node(.empty, 7, .empty)
let node6 = BinarySearchTree.node(.empty, 6, node7)
let node5 = BinarySearchTree.node(.empty, 5, node6)
let node4 = BinarySearchTree.node(node3, 4, node5)

var sortedArray: [Int] = []
node4.traverseInOrder { (value:Int) in
    sortedArray.append(value)
}
print(sortedArray) // [1, 2, 3, 4, 5, 6, 7]

let bst = BinarySearchTree.sortedArrayToBST(sortedArray, start: 0, end: sortedArray.count - 1)
print(bst) // (((. <- 1 -> .) <- 2 -> (. <- 3 -> .)) <- 4 -> ((. <- 5 -> .) <- 6 -> (. <- 7 -> .)))