Add One Row to Tree

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem

Add a row of nodes with value val at depth depth. Existing subtrees become children of new nodes.

Solutions

# Python
from collections import deque

def addOneRow(root, val, depth):
    if depth == 1:
        new_root = TreeNode(val)
        new_root.left = root
        return new_root
    q = deque([root])
    cur_depth = 1
    while q:
        if cur_depth == depth - 1:
            for node in q:
                new_left = TreeNode(val)
                new_right = TreeNode(val)
                new_left.left = node.left
                new_right.right = node.right
                node.left = new_left
                node.right = new_right
            break
        cur_depth += 1
        for _ in range(len(q)):
            node = q.popleft()
            if node.left:  q.append(node.left)
            if node.right: q.append(node.right)
    return root
// C++
TreeNode* addOneRow(TreeNode* root, int val, int depth) {
    if (depth == 1) { auto r = new TreeNode(val); r->left = root; return r; }
    queue<TreeNode*> q; q.push(root); int d = 1;
    while (d < depth - 1) {
        int sz = q.size();
        while (sz--) { auto node = q.front(); q.pop(); if(node->left)q.push(node->left); if(node->right)q.push(node->right); }
        d++;
    }
    while (!q.empty()) {
        auto node = q.front(); q.pop();
        auto nl = new TreeNode(val); nl->left = node->left; node->left = nl;
        auto nr = new TreeNode(val); nr->right = node->right; node->right = nr;
    }
    return root;
}
// Java
public TreeNode addOneRow(TreeNode root, int val, int depth) {
    if (depth == 1) { TreeNode r = new TreeNode(val); r.left = root; return r; }
    Queue<TreeNode> q = new LinkedList<>(); q.add(root); int d = 1;
    while (d < depth - 1) {
        int sz = q.size();
        while (sz-- > 0) { TreeNode n = q.poll(); if(n.left!=null)q.add(n.left); if(n.right!=null)q.add(n.right); }
        d++;
    }
    for (TreeNode node : q) {
        TreeNode nl = new TreeNode(val); nl.left = node.left; node.left = nl;
        TreeNode nr = new TreeNode(val); nr.right = node.right; node.right = nr;
    }
    return root;
}
// JavaScript
function addOneRow(root, val, depth) {
    if (depth === 1) { const r = new TreeNode(val); r.left = root; return r; }
    let queue = [root], d = 1;
    while (d < depth - 1) { queue = queue.flatMap(n => [n.left, n.right].filter(Boolean)); d++; }
    for (const node of queue) {
        const nl = new TreeNode(val); nl.left = node.left; node.left = nl;
        const nr = new TreeNode(val); nr.right = node.right; node.right = nr;
    }
    return root;
}

Complexity

  • Time: O(n)
  • Space: O(n)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro