BinaryTree.swift

import Foundation

class Node<T: Comparable> {
    let value: T
    var left: Node?
    var right: Node?

    init(value: T, left: Node? = nil, right: Node? = nil) {
        self.value = value
        self.left = left
        self.right = right
    }

    /// Creates a Binary Search Tree rooted at the returned Node,
    /// with the values in `values`. Returns `nil` if values.count == 0
    static func binarySearchTree(fromValues values: [T]) -> Node? {
        guard values.count > 0 else {
            return nil
        }

        let root = Node(value: values.first!)

        guard values.count > 1 else {
            return root
        }

        for value in values[1 ..< values.count] {
            insert(value, into: root)
        }

        return root
    }

    /// Inserts a new Node with value of `value` into the binary tree
    /// rooted at `root`
    static func insert(v: T, into root: Node) {

        if v < root.value {
            if root.left != nil {
                insert(v, into: root.left!)
            } else {
                root.left = Node(value: v)
            }
        } else {
            if root.right != nil {
                insert(v, into: root.right!)
            } else {
                root.right = Node(value: v)
            }
        }

    }

    /// Returns an integer representing the distance from a parent node to
    /// a child node, or `nil` if `node` is not a child of `parent`. Returns
    /// 0 if root.value == node.value
    static func distance(fromParent parent: Node, toChild child: Node) -> Int? {
        if child.value == parent.value {
            return 0
        }

        guard let newParent = child.value < parent.value ? parent.left : parent.right else {
            return nil
        }

        guard let distance = distance(fromParent: newParent, toChild: child) else {
            return nil
        }

        return distance + 1
    }

    /// Returns an integer representing the distance between two nodes, or `nil`
    /// if one or both nodes do not exist as children of `root`.
    static func distance(`in` root: Node, from node1: Node, to node2: Node) -> Int? {
        guard let commonParent = lowestCommonParent(in: root, of: node1, and: node2) else {
            return nil
        }

        guard let d1 = distance(fromParent: commonParent, toChild: node1),
            let d2 = distance(fromParent: commonParent, toChild: node2) else {
                return nil
        }

        let totalDistance = d1 + d2
        return totalDistance
    }

    /// Returns the deepest node in the tree that is a parent of both `node1` and `node2`, or
    /// `nil` if `node1` and `node2` do not share a common parent
    static func lowestCommonParent(`in` root: Node, of node1: Node, and node2: Node) -> Node? {
        // If n1 < root <= n2, root is in between, and therefore the lowest common parent
        if node1.value < root.value && root.value <= node2.value {
            return root
        }

        // Same idea; if n2 < root <= n1, root is in between, and therefore the lowest common parent
        if node2.value < root.value && root.value <= node1.value {
            return root
        }

        let distanceToNode1 = distance(fromParent: root, toChild: node1)
        let distanceToNode2 = distance(fromParent: root, toChild: node2)

        guard let d1 = distanceToNode1, d2 = distanceToNode2 else {
            return nil
        }

        if d1 == 0 || d2 == 0 {
            return root
        }

        // We now know that both node1 & node2 are to the same side of `root`.
        guard let newRoot = node1.value < root.value ? root.left : root.right else {
            return nil
        }

        return lowestCommonParent(in: newRoot, of: node1, and: node2)
    }

    /// Returns the child node of `searchRoot` whose value == `value`.
    static func find(v: T, `in` searchRoot: Node) -> Node? {
        if v == searchRoot.value {
            return searchRoot
        }

        guard let newSearchRoot = v < searchRoot.value ? searchRoot.left : searchRoot.right else {
            return nil
        }

        return find(v, in: newSearchRoot)
    }


}

// MARK: - CustomStringConvertible

extension Node: CustomStringConvertible {
    var description: String {
        return "Node«\(value); left: \(left); right: \(right)»"
    }
}


// MARK: - Main

let root = Node<Int>.binarySearchTree(fromValues: [5, 3, 8, 1, 4, 2, 9, 6])!
print("root: \(root)")

let n3 = Node.find(3, in: root)!
let distanceFromRootToN3 = Node.distance(fromParent: root, toChild: n3)!
print("n3: \(n3)")
print("distanceFromRootToN3: \(distanceFromRootToN3)")

let n5 = Node.find(5, in: root)!
let distanceFromRootToN5 = Node.distance(fromParent: root, toChild: n5)!
print("distanceFromRootToN5: \(distanceFromRootToN5)")
print("root===n5: \(root === n5)")

let n9 = Node.find(9, in: root)!
let distanceFromRootToN9 = Node.distance(fromParent: root, toChild: n9)!
print("distanceFromRootToN9: \(distanceFromRootToN9)")

let parentOf3And9 = Node.lowestCommonParent(in: root, of: n3, and: n9)
print("root===parentOf3And9: \(root === parentOf3And9)")
print("parentOf3And9: \(parentOf3And9)")

let distanceFrom3To9 = Node.distance(in: root, from: n3, to: n9)
print("distanceFrom3To9: \(distanceFrom3To9)")