[Tree] Refactor solution to Binary Tree Level Order Traversal

Author: soapyigu

Tree/BinaryTreeLevelOrderTraversal.swift

@@ -17,36 +17,25 @@
  */
 
 class BinaryTreeLevelOrderTraversal {
-    func levelOrder(root: TreeNode?) -> [[Int]] {
+    func levelOrder(_ root: TreeNode?) -> [[Int]] {
         var res = [[Int]]()
-        var queue = [TreeNode]()
 
-        if let root = root {
-            queue.append(root)
+        guard let root = root else {
+            return res
         }
 
-        while queue.count > 0 {
-            var size = queue.count
-            var level = [Int]()
+        var currentLevel = [root]
+        
+        while !currentLevel.isEmpty {
+            let currentLevelVals = currentLevel.map { $0.val }
+            
+            // add current level vals
+            res.append(currentLevelVals)
 
-            for _ in 1...size {
-                let node = queue[0]
-                queue.removeAtIndex(0)
-                
-                // add val
-                level.append(node.val)
-
-                // add TreeNodes in next level
-                if let left = node.left {
-                    queue.append(left)
-                }
-                if let right = node.right {
-                    queue.append(right)
-                }
-            }
-            res.append(level)
+            // add next level nodes
+            currentLevel = currentLevel.flatMap { [$0.left, $0.right] }.compactMap { $0 }
         }
 
         return res
     }
-}
+}