Leetcode 102: Binary tree level-order traversal in Swift

Question

102. Binary Tree Level Order Traversal

Given a binary tree, return the level order traversal of its nodes' values. (ie, from left to right, level by level).

For example:

Given binary tree [3,9,20,null,null,15,7],

    3
   / \
  9  20
    /  \
   15   7

return its level order traversal as:

[
  [3],
  [9,20],
  [15,7]
]

This is my solution in Swift:

func levelOrder(_ root: TreeNode?) -> [[Int]] {
        guard let root = root else{
            return []
        }
        let nodes = recursiveVisit(root)
        let capacity = nodes.reduce([Int]()) {
            if $0.contains($1.0) == false {
                return $0 + [$1.0]
            }
            return $0
        }.count
        return nodes.reduce([[Int]](repeating: [], count: capacity), { 
          var tmp = $0
            tmp[$1.0].append($1.1)
            return tmp
        })
    }


func recursiveVisit(_ node: TreeNode?) -> [(Int, Int)]{
        //  [(Int,  Int)]
        //  depth,      node.val

        guard let node = node else{
            return []
        }

        var nodes = [(Int, Int)]()
        nodes.append((0, node.val))

        let lhs = recursiveVisit(node.left).map {
            return ($0.0 + 1, $0.1)
        }

        let rhs = recursiveVisit(node.right).map {
            return ($0.0 + 1, $0.1)
        }
        nodes.append(contentsOf: lhs)
        nodes.append(contentsOf: rhs)
        return nodes
    }

The solution is very intuitive: collect the nodes' value and depth, then reduce to the answer.

It is not as well designed as the Java Solution, though.

My question is mainly on the Swift syntax. Is there any other way to do it more lean? Combine the two reduce() calls to one, for example?

Answer 1

Simplifying your code

Let's start here:

let capacity = nodes.reduce([Int]()) {
    if $0.contains($1.0) == false {
        return $0 + [$1.0]
    }
    return $0
}.count 

The test

if $0.contains($1.0) == false

is shorter (and – in my opinion – better) written as

if !$0.contains($1.0)

Each invocation of the closure creates a new array. In such cases it is more efficient to use reduce(into:) with a closure that updates the array:

let capacity = nodes.reduce(into: [Int]()) {
    if !$0.contains($1.0) {
        $0.append($1.0)
    }
}.count

But what this code actually does is to determine the number of distinct levels in the nodes array. That is simpler and more efficiently done with a set:

let capacity = Set(nodes.map { $0.0 }).count

Here

return nodes.reduce([[Int]](repeating: [], count: capacity), {
    var tmp = $0
    tmp[$1.0].append($1.1)
    return tmp
})

a new nested array is created with each invocation of the closure. Again, this can be improved with reduce(into:), so that the first function becomes

func levelOrder(_ root: TreeNode?) -> [[Int]] {
    guard let root = root else {
        return []
    }
    let nodes = recursiveVisit(root)
    let capacity = Set(nodes.map { $0.0 }).count
    return nodes.reduce(into: [[Int]](repeating: [], count: capacity), {
        $0[$1.0].append($1.1)
    })
}

I have only minor suggestions for the second function:

var nodes = [(Int, Int)]()
nodes.append((0, node.val))

can be combined to

var nodes = [(0, node.val)]

The results of the recursive calls can be appended to the nodes array directly:

nodes += recursiveVisit(node.left).map {
    return ($0.0 + 1, $0.1)
}
nodes += recursiveVisit(node.right).map {
    return ($0.0 + 1, $0.1)
}

And if you use tuple labels then the code becomes better readable (and almost self-documenting):

func recursiveVisit(_ node: TreeNode?) -> [(level: Int, value: Int)] {

    guard let node = node else {
        return []
    }
    var nodes = [(level: 0, value: node.val)]
    nodes += recursiveVisit(node.left).map {
        return (level: $0.level + 1, value: $0.value)
    }
    nodes += recursiveVisit(node.right).map {
        return (level: $0.level + 1, value: $0.value)
    }
    return nodes
}

This can be used in the first function as well. As an alternative, define a custom struct with “level” and “label” members.

Alternative #1 – A nested function

It is not necessary to create an array of (level, value) tuples first. You can add a value to the sublist on the current level, or append a new level while traversing the tree recursively. With a nested function you don't even have to pass the array around:

func levelOrder(_ root: TreeNode?) -> [[Int]] {
    var levels: [[Int]] = []

    func recursiveVisit(_ node: TreeNode?, level: Int) {
        guard let node = node else {
            return
        }
        if level < levels.count {
            levels[level].append(node.val)
        } else {
            levels.append([node.val])
        }
        recursiveVisit(node.left, level: level + 1)
        recursiveVisit(node.right, level: level + 1)
    }

    recursiveVisit(root, level: 0)
    return levels
}

Alternative #2 – Iteration instead of recursion

The referenced Java solution solves the task with iteration, and a list containing all nodes on the current level. Of course that can be done in Swift as well. Here is an example of a quite compact implementation:

func levelOrder(_ root: TreeNode?) -> [[Int]] {
    guard let root = root else {
        return []
    }

    var wrapList: [[Int]] = []
    var queue = [root] // First level

    while !queue.isEmpty {
        // All values of nodes on the current level:
        wrapList.append(queue.map { $0.val })
        // Replace queue by list of all nodes on the next level:
        queue = queue.flatMap { [$0.left, $0.right ] }.compactMap { $0 }
    }

    return wrapList
}