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
}