BinarySearchTree

BinarySearchTree.md

public class BinaryNode<Element> { 
  public var value: Element
  public var leftChild: BinaryNode?
  public var rightChild: BinaryNode?
  
  public init(value: Element) {
    self.value = value
  }
}

public struct BinarySearchTree<Element: Comparable> {
  public private(set) var root: BinaryNode<Element>?
  public init() {}
}

extension BinarySearchTree {
  public mutating func insert(_ value: Element) {
    root = insert(from: root, value: value)
  }
  
  /// O(log n)
  private func insert(
    from node: BinaryNode<Element>?,
    value: Element
  ) -> BinaryNode<Element> {
    // base case
    guard let node = node else {
      return BinaryNode(value: value)
    }
    if value < node.value {
      node.leftChild = insert(from: node.leftChild, value: value)
    } else {
      node.rightChild = insert(from: node.rightChild, value: value)
    }
    return node
  }
}

extension BinarySearchTree {
  // O(log n)
  public func contains(_ value: Element) -> Bool {
    var current = root
    while let node = current {
      if node.value == value {
        return true
      }
      if value < node.value {
        current = node.leftChild
      } else {
        current = node.rightChild
      }
    }
    return false
  }
}

private extension BinaryNode {
  var min: BinaryNode {
    return leftChild?.min ?? self
  }
}

extension BinarySearchTree {
  public mutating func remove(_ value: Element) {
    root = remove(node: root, value: value)
  }
  
  private func remove(
    node: BinaryNode<Element>?,
    value: Element
  ) -> BinaryNode<Element>? {
    // (base case) leaf node
    guard let node = node else { return nil }
    if value == node.value {
      // (base case) leaf node
      if node.leftChild == nil && node.rightChild == nil {
        return nil
      }
      // (base case) return node.leftChild to reconnect the right subtree.
      if node.leftChild == nil {
        return node.rightChild
      }
      // (base case) return node.rightChild to reconnect the left subtree.
      if node.rightChild == nil {
        return node.leftChild
      }
      // (base case) the node to be removed has both a left and right child.
      node.value = node.rightChild!.min.value
      node.rightChild = remove(node: node.rightChild, value: node.value)
    } else if value < node.value {
      node.leftChild = remove(node: node.leftChild, value: value)
    } else {
      node.rightChild = remove(node: node.rightChild, value: value)
    }
    // (base case) leaf node
    return node
  }
}