AverageOfLevelsInBST.swift

public class TreeNode {
  public var val: Int
  public var left: TreeNode?
  public var right: TreeNode?
  public init() { self.val = 0; self.left = nil; self.right = nil; }
  public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
  public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
    self.val = val
    self.left = left
    self.right = right
  }
}

extension TreeNode {
    var height: Int {
        return 1 + Swift.max(leftHeight, rightHeight)
    }
    
    var leftHeight: Int {
        return (left?.height) ?? 0
    }
    
    var rightHeight: Int {
        return (right?.height) ?? 0
    }
    
    var isLeaf: Bool {
        return left == nil && right == nil
    }
    
    var validSubtrees: [TreeNode] {
        var subtrees = [TreeNode]()
        if let left = left {
            subtrees.append(left)
        }
        
        if let right = right {
            subtrees.append(right)
        }
        
        return subtrees
    }
}

extension Collection where Element: BinaryInteger {
    var sum: Element {
        return reduce(0, +)
    }
    
    var average: Double {
        return Double(sum) / Double(count)
    }
}

class Solution {
    func averageOfLevels(_ root: TreeNode?) -> [Double] {
        guard let root = root else {
            return []
        }
        
        return randomSolution(root)
    }
    
    func randomSolution(_ root: TreeNode) -> [Double] {
        print("Running random solution: ")
        
        let randomSolutionClosure:() -> [Double] = [{
            print("recursive-functional")
            return self.averageOfLevelsRecursive_Functional([root])
        }, {
            print("recursive-imperative")
            return self.averageOfLevelsRecursive_Imperative([root])
        }, {
            print("iterative-functional")
            return self.averageOfLevelsIterative_Functional(root)
        }, {
            print("iterative-imperative")
            return self.averageOfLevelsIterative_Imperative(root)
        }].randomElement()!
        
        return randomSolutionClosure()
    }
    
    func averageOfLevelsRecursive_Functional(_ queue: [TreeNode]) -> [Double] {
        guard queue.isEmpty == false else {
            return []
        }
        
        let reduceSeed = (values: [Int](), nextLevel: [TreeNode]())
        let (values, nextLevel) = queue.reduce(reduceSeed) { res, node in
            return (res.values + [node.val], res.nextLevel + node.validSubtrees)
        }
        
        return [values.average] + averageOfLevelsRecursive_Functional(nextLevel)
    }
    
    func averageOfLevelsRecursive_Imperative(_ queue: [TreeNode]) -> [Double] {
        guard queue.isEmpty == false else {
            return []
        }

        var values = [Int]()
        var nextLevel = [TreeNode]()
        
        queue.forEach { node in 
            values.append(node.val)

            nextLevel.append(contentsOf: node.validSubtrees)
        }
        
        return [values.average] + averageOfLevelsRecursive_Imperative(nextLevel)
    }
    
    func averageOfLevelsIterative_Functional(_ node: TreeNode) -> [Double] {
        // Unavoidable state.
        var queue = [node]
        
        // Need to compute height to know how many times we need to reduce.
        // Expensive. :/
        
        return (1...node.height).reduce([Double]()) { res, level in
            let newResult = (queue.map { $0.val }).average
            queue = queue.flatMap { $0.validSubtrees }
            return res + [newResult]
        }
    }
    
    func averageOfLevelsIterative_Imperative(_ node: TreeNode) -> [Double] {
        var results = [Double]()
        var queue = [node]
        
        while !queue.isEmpty {
            results.append(queue.map { $0.val }.average)
            queue = queue.flatMap { $0.validSubtrees }
        }

        return results
    }
}