LeetCode-in-Swift

114. Flatten Binary Tree to Linked List

Medium

Given the root of a binary tree, flatten the tree into a “linked list”:

Example 1:

Input: root = [1,2,5,3,4,null,6]

Output: [1,null,2,null,3,null,4,null,5,null,6]

Example 2:

Input: root = []

Output: []

Example 3:

Input: root = [0]

Output: [0]

Constraints:

Follow up: Can you flatten the tree in-place (with O(1) extra space)?

Solution

/**
 * Definition for a binary tree node.
 * 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
 *     }
 * }
 */
class Solution {
    func flatten(_ root: TreeNode?) {
        var values: [Int] = []
        getAllValues(root, &values)
        let linkedList = TreeNode()
        var currLink: TreeNode? = linkedList
        for index in 0..<values.count {
            currLink?.right = TreeNode(values[index])
            currLink = currLink?.right
        }
        root?.right = linkedList.right?.right
        root?.left = nil
    }

    private func getAllValues(_ curr: TreeNode?, _ values: inout [Int]) {
        guard let curr = curr else { return }
        values.append(curr.val)
        if let left = curr.left {
            getAllValues(left, &values)
        }
        if let right = curr.right {
            getAllValues(right, &values)
        }
    }
}