623.add-one-row-to-tree.md

623. Add One Row to Tree

BFS

We can use BFS to traverse the tree, find the k - 1 depth nodes (not k depth), and add the new nodes to the left and right of the nodes. And the original left and right nodes will be the left and right nodes of the new nodes.

     A      B       C     // depth = k - 1
   /          \   /   \
  1            2 3     4  // depth = k
A.left = 1
B.right = 2
C.left = 3
C.right = 4

// Add one row
    A       B               C     // depth = k - 1
   /  \   /   \           /   \
  K    K K     K         K     K  // depth = k
 /               \      /       \
1                 2    3         4

A.left = K
A.left.left = 1

B.left = K
B.left.right = 2

C.left = K
C.left.left = 3
C.right = K
C.right.right = 4

fun addOneRow(root: TreeNode?, `val`: Int, depth: Int): TreeNode? {
    if (root == null) return null
    if (depth == 1) {
        val newRoot = TreeNode(`val`)
        newRoot.left = root
        return newRoot
    }
    val queue = ArrayDeque<TreeNode>()
    var currentDepth = 1
    queue.addLast(root)
    while (queue.isNotEmpty()) {
        val size = queue.size
        for (i in 0 until size) {
            val node = queue.removeFirst()
            if (currentDepth == depth - 1) {
                val newLeft = TreeNode(`val`)
                newLeft.left = node.left
                node.left = newLeft

                val newRight = TreeNode(`val`)
                newRight.right = node.right
                node.right = newRight
            } else {
                if (node.left != null) queue.addLast(node.left)
                if (node.right != null) queue.addLast(node.right)
            }
        }
        currentDepth++
    }
    return root
}

• Time Complexity: O(n), where n is the number of nodes in the tree.
• Space Complexity: O(n).

DFS

We can also implement in DFS approach, the idea is the same as BFS:

fun addOneRow(root: TreeNode?, `val`: Int, depth: Int): TreeNode? {
    if (root == null) return null
    if (depth == 1) {
        val newRoot = TreeNode(`val`)
        newRoot.left = root
        return newRoot
    } else {
        dfs(root, `val`, 1, depth)
        return root
    }
}

// No return value, we modify the tree in place.
private fun dfs(root: TreeNode?, `val`: Int, currentDepth: Int, depth: Int) {
    if (root == null) return
    if (currentDepth == depth - 1) {
        val newLeft = TreeNode(`val`)
        newLeft.left = root.left
        root.left = newLeft

        val newRight = TreeNode(`val`)
        newRight.right = root.right
        root.right = newRight
        return 
    }
    dfs(root.left, `val`, currentDepth + 1, depth)
    dfs(root.right, `val`, currentDepth + 1, depth)
}

// Or equivalently, `dfs()` return the new node.
fun addOneRow(root: TreeNode?, `val`: Int, depth: Int): TreeNode? {
    if (root == null) return null
    return dfs(root, `val`, 1, depth)
}

private fun dfs(root: TreeNode?, `val`: Int, currentDepth: Int, depth: Int): TreeNode? {
    if (root == null) return null
    if (depth == 1) {
        val newRoot = TreeNode(`val`)
        newRoot.left = root
        return newRoot
    }
    if (currentDepth == depth - 1) {
        val newLeft = TreeNode(`val`)
        newLeft.left = root.left
        root.left = newLeft

        val newRight = TreeNode(`val`)
        newRight.right = root.right
        root.right = newRight
        return root
    }
    root.left = dfs(root.left, `val`, currentDepth + 1, depth)
    root.right = dfs(root.right, `val`, currentDepth + 1, depth)
    return root
}

// Or equivalently, we implement without helper function.
fun addOneRow(root: TreeNode?, `val`: Int, depth: Int): TreeNode? {
    if (root == null) return null
    if (depth == 1) {
        val newRoot = TreeNode(`val`)
        newRoot.left = root
        return newRoot
    }

    if (depth == 2) {
        val newLeft = TreeNode(`val`)
        newLeft.left = root.left
        root.left = newLeft

        val newRight = TreeNode(`val`)
        newRight.right = root.right
        root.right = newRight

        return root
    }

    root.left = addOneRow(root.left, `val`, depth - 1)
    root.right = addOneRow(root.right, `val`, depth - 1)
    return root
}