Swift: arrayToTree() where array contains int and nil

Question

I am learning tree on Leetcode. Need to prepare the testing data.

It is easy to convert the array to an organized node, where its elements are integers.

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

Here is my code:

extension Array where Element == Int{

func arrayToTree() -> TreeNode{
    var nodes = [TreeNode]()
    for num in 0..<self.count{
         nodes.append(TreeNode(self[num]))
    }

    var i = 0
    repeat{
        nodes[i].left = nodes[2 * i + 1]
        if self.count > 2 * i + 2 {
             nodes[i].right = nodes[2 * i + 2]
        }
        i+=1
    }while i < (self.count)/2

    return nodes.first!
  }
}

When the last level contains some nils , such as [3, 9, 20, nil, nil, 15, 7]

Here is the code:

extension Array where Element == Int?{

    func arrayToTree() -> TreeNode?{

        guard self.count > 0 else{
            return nil
        }
        var nodes = [TreeNode?]()
        for num in 0..<self.count{
            if let num = self[num]{
                nodes.append(TreeNode(num))
            }
            else{
                nodes.append(nil)
            }
        }
        var i = 0
        repeat {
            nodes[i]?.left = nodes[2 * i + 1]
            if self.count > 2 * i + 2 {
                nodes[i]?.right = nodes[2 * i + 2]
            }
            i += 1
        } while i < (self.count) / 2

        return nodes.first!
    }


}

How can I refactor this?

Such as combine them together using Swift Generic with some Protocol.

Answer 1

Simplifying func arrayToTree()

This

   var nodes = [TreeNode?]()
   for num in 0..<self.count{
       if let num = self[num]{
           nodes.append(TreeNode(num))
       }
       else{
           nodes.append(nil)
       }
   }

creates a new array by mapping each element in self (an optional Int) to a new element (an optional TreeNode). That can be simplified to

let nodes = self.map { $0.map { TreeNode($0) } }

where the outer Array.map maps the given array to a new array, and the inner Optional.map maps an optional Int to an optional TreeNode.

The

   var i = 0
   repeat {
       // ...
       i += 1
   } while i < (self.count) / 2

loop can be simplified to

for i in 0..<self.count/2 {
    // ...
}

The forced unwrapping

   return nodes.first!

cannot crash – the nodes array cannot be empty at this point. I would still suggest to avoid it since later code changes might break the logic. It also makes it easier for future maintainers of the code to verify its correctness.

Actually the preceding code just results in an empty nodes array if the given list is empty. Therefore we can remove the initial guard and replace it by

guard let first = nodes.first else {
    return nil
}
return first

at the end of the method. This can be further shortened to

return nodes.first.flatMap { $0 }

Putting it together, the function would look like this:

func arrayToTree() -> TreeNode? {

    let nodes = self.map { $0.map { TreeNode($0) } }
    for i in 0..<self.count/2 {
        nodes[i]?.left = nodes[2 * i + 1]
        if self.count > 2 * i + 2 {
            nodes[i]?.right = nodes[2 * i + 2]
        }
    }

    return nodes.first.flatMap { $0 }
}

An alternative implementation

What you have is a way to create a TreeNode, and that is what initializers methods are for. Therefore I would put the code in a

public convenience init?(values: [Int?])

of the TreeNode class instead of an Array extension method. The usage would be

let list = [3, 9, 20, nil, nil, 15, 7]
if let tree = TreeNode(values: list) {
    // ...
}

And the task calls for a recursive implementation:

public convenience init?(values: [Int?], offset: Int = 0) {
    guard offset < values.count, let value = values[offset] else {
        return nil
    }
    self.init(value)
    self.left = TreeNode(values: values, offset: 2 * offset + 1)
    self.right = TreeNode(values: values, offset: 2 * offset + 2)
}

Making it generic

With the above changes it is now easy to replace Int by an arbitrary value type:

public class TreeNode<T> {
    public var val: T
    public var left: TreeNode?
    public var right: TreeNode?
    public init(_ val: T) {
        self.val = val
        self.left = nil
        self.right = nil
    }

    public convenience init?(values: [T?], offset: Int = 0) {
        guard offset < values.count, let value = values[offset] else {
            return nil
        }
        self.init(value)
        self.left = TreeNode(values: values, offset: 2 * offset + 1)
        self.right = TreeNode(values: values, offset: 2 * offset + 2)
    }
}