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
    }
    
    public convenience init?(array: [Int?]) {
        guard !array.isEmpty else { return nil }
        
        var index = 0
        let root = TreeNode(array[index]!)
        index += 1
        
        var queue: [TreeNode] = [root]
        
        while index < array.count {
            let current = queue.removeFirst()
            
            if index < array.count, let leftValue = array[index] {
                current.left = TreeNode(leftValue)
                queue.append(current.left!)
            }
            index += 1
            
            if index < array.count, let rightValue = array[index] {
                current.right = TreeNode(rightValue)
                queue.append(current.right!)
            }
            index += 1
        }
        
        self.init(root.val, root.left, root.right)
    }
}

extension TreeNode {
    public func printTree() {
        printTreeHelper(self, "", true)
    }
    
    private func printTreeHelper(_ node: TreeNode?, _ prefix: String, _ isLeft: Bool) {
        guard let node = node else { return }
        
        let arrow = isLeft ? "┣━━" : "┗━━"
        let branch = isLeft ? "┃  " : "   "
        
        print(prefix + arrow + " " + "\(node.val)")
        
        printTreeHelper(node.left, prefix + branch, true)
        printTreeHelper(node.right, prefix + branch, false)
    }
}

 func levelOrder(_ root: TreeNode?) -> [[Int]] {
        var result = [[Int]]()
        
        guard let root else { return result }
        
        var currentLevel = [root]
        
        while !currentLevel.isEmpty {
            let currentLevelValues = currentLevel.map(\.val)
            
            // add current level values
            result.append(currentLevelValues)
            
            // add next level nodes
            currentLevel = currentLevel.flatMap {[$0.left, $0.right]}.compactMap { $0 }
        }
        
        return result
    }
    
    let tree1 = TreeNode(array: [3,9,20,nil,nil,15,7])
    print(levelOrder(tree1)) // [[3], [9, 20], [15, 7]]

    let tree2 = TreeNode(array: [1])
    print(levelOrder(tree2)) // [[1]]

    let tree3 = TreeNode(array: [])
    print(levelOrder(tree3)) // []
}

func levelOrder(_ root: TreeNode?) -> [[Int]] {
    var levels: [[Int]] = []
    
    guard let root else { return levels }
    
    func helper(_ node: TreeNode, _ level: Int) {
        if levels.count == level {
            levels.append([])
        }
        
        levels[level].append(node.val)
        
        if let leftNode = node.left {
            helper(leftNode, level + 1)
        }
        
        if let rightNode = node.right {
            helper(rightNode, level + 1)
        }
    }
    
    helper(root, 0)
    
    return levels
}

let tree1 = TreeNode(array: [3,9,20,nil,nil,15,7])
print(levelOrder(tree1)) // [[3], [9, 20], [15, 7]]

let tree2 = TreeNode(array: [1])
print(levelOrder(tree2)) // [[1]]

let tree3 = TreeNode(array: [])
print(levelOrder(tree3)) // []