import UIKit

struct Queue<Element> {
    
    var elements: [Element] = []
    
    mutating func enqueue(_ value: Element) -> Bool {
        elements.append(value)
        return true
    }
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    mutating func dequeue() -> Element? {
        return isEmpty ? nil : elements.removeFirst()
    }
}


class TreeNode<T> {
    
    var value: T
    var children: [TreeNode] = []
    
    init(_ value: T) {
        self.value = value
    }
    
    func add(_ child: TreeNode) {
        self.children.append(child)
    }
}

extension TreeNode {
    
    func forEachLevelOrder(_ visit: (TreeNode) -> Void) {
        visit(self)                                             //call root node
        
        var queue = Queue<TreeNode>()
        children.forEach {
            queue.enqueue($0)
        }
        while let node = queue.dequeue() {
            visit(node)
            node.children.forEach {
                queue.enqueue($0)
            }
        }
    }
}

/*
           Beverages
           /       \
        Hot         Cold
       /   \        /   \
    Tea   Coffee   Soda  Milk

*/

let beverages = TreeNode<String>("Beverages")
let hot = TreeNode<String>("Hot")
let cold = TreeNode<String>("Cold")
beverages.add(hot)
beverages.add(cold)

let tea = TreeNode<String>("Tea")
let coffee = TreeNode<String>("Coffee")
hot.add(tea)
hot.add(coffee)

let soda = TreeNode<String>("Soda")
let milk = TreeNode<String>("Milk")
cold.add(soda)
cold.add(milk)


beverages.forEachLevelOrder {
    print($0.value)     // level order traversal
}

/* 
Output printed:
    Beverages
    Hot
    Cold
    Tea
    Coffee
    Soda
    Milk
*/
