BinarySearchTree

BinarySearchTree.swift

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

extension BinarySearchTree {
  
  var count: Int {
    switch self {
    case .empty: return 0
    case .leaf: return 1
    case let .node(left, _, right): return left.count + 1 + right.count
    }
  }
  
  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)
    }
  }
  
  func insert(_ newValue: T) -> BinarySearchTree {
    switch self {
    case .empty:
      return .leaf(newValue)
      
    case .leaf(let value):
      return newValue < value ?
        .node(.leaf(newValue), value, .empty) :
        .node(.empty, value, .leaf(newValue))
      
    case .node(let left, let value, let right):
      return newValue < value ?
        .node(left.insert(newValue), value, right) :
        .node(left, value, right.insert(newValue))
    }
  }
  
  func search( _ value: T) -> BinarySearchTree? {
    switch self {
    case .empty: return nil
    case let .leaf(leaf): return leaf == value ? self : nil
    case let .node(right, leaf, left):
      switch value {
      case leaf: return self
      case value where value > leaf: return left.search(value)
      case value where value < leaf: return right.search(value)
      default: return nil
      }
    }
  }
}