Add One Row to Tree

medium tree dfs bfs

Problem

Given the root of a binary tree, add a row of nodes with value v at depth d. The old left child of every depth-(d−1) node becomes the new left node's left subtree; similarly for right. If d == 1, make a brand new root.

Inputroot = [4,2,6,3,1,5], v = 1, d = 2

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

At depth 1 every node gets two new value-1 children spliced above its old children.

tree (level-order)

def add_one_row(root, v, d):
    if d == 1:
        return TreeNode(v, root, None)
    def dfs(node, depth):
        if not node:
            return
        if depth == d - 1:
            node.left  = TreeNode(v, node.left, None)
            node.right = TreeNode(v, None, node.right)
        else:
            dfs(node.left,  depth + 1)
            dfs(node.right, depth + 1)
    dfs(root, 1)
    return root

function addOneRow(root, v, d) {
  if (d === 1) return { val: v, left: root, right: null };
  function dfs(node, depth) {
    if (!node) return;
    if (depth === d - 1) {
      node.left  = { val: v, left: node.left,  right: null };
      node.right = { val: v, left: null,       right: node.right };
    } else {
      dfs(node.left,  depth + 1);
      dfs(node.right, depth + 1);
    }
  }
  dfs(root, 1);
  return root;
}

class Solution {
    public TreeNode addOneRow(TreeNode root, int v, int d) {
        if (d == 1) return new TreeNode(v, root, null);
        dfs(root, 1, v, d);
        return root;
    }
    private void dfs(TreeNode node, int depth, int v, int d) {
        if (node == null) return;
        if (depth == d - 1) {
            node.left  = new TreeNode(v, node.left,  null);
            node.right = new TreeNode(v, null, node.right);
        } else {
            dfs(node.left,  depth + 1, v, d);
            dfs(node.right, depth + 1, v, d);
        }
    }
}

TreeNode* addOneRow(TreeNode* root, int v, int d) {
    if (d == 1) return new TreeNode(v, root, nullptr);
    function<void(TreeNode*, int)> dfs = [&](TreeNode* node, int depth) {
        if (!node) return;
        if (depth == d - 1) {
            node->left  = new TreeNode(v, node->left,  nullptr);
            node->right = new TreeNode(v, nullptr, node->right);
        } else {
            dfs(node->left,  depth + 1);
            dfs(node->right, depth + 1);
        }
    };
    dfs(root, 1);
    return root;
}

Explanation

We want to slip a whole new row of value v into the tree at depth d. The key insight is that the new nodes don't go at depth d by accident — they get attached as children of the nodes sitting at depth d - 1.

So we walk down with a DFS, tracking the current depth. When we reach a node at depth d - 1, we stop recursing and do the splice: its old left child becomes the left child of a new value-v node, and its old right child becomes the right child of another new value-v node.

This keeps the old subtrees attached, just one level lower, with the new row wedged in between.

There is one special case: if d == 1, there is no depth-0 parent to attach to, so we make a brand new root holding v with the entire old tree as its left child.

Example: tree [4,2,6,...] with v = 1, d = 2. Node 4 is at depth 1, so we give it two new 1 nodes; 4's old left child 2 hangs under the left 1, and old right child 6 hangs under the right 1.

Time: O(n) Space: O(h) recursion stack