[LeetCode] 623. Add One Row to Tree

题目描述

Given the root of a binary tree, then value v and depth d, you need to add a row of nodes with value v at the given depth d. The root node is at depth 1.

The adding rule is: given a positive integer depth d, for each NOT null tree nodes N in depth d-1, create two tree nodes with value v as N’s left subtree root and right subtree root. And N’s original left subtree should be the left subtree of the new left subtree root, its original right subtree should be the right subtree of the new right subtree root. If depth d is 1 that means there is no depth d-1 at all, then create a tree node with value v as the new root of the whole original tree, and the original tree is the new root’s left subtree.

Example 1:

Input: 
A binary tree as following:
       4
     /   \
    2     6
   / \   / 
  3   1 5   

v = 1

d = 2

Output: 
       4
      / \
     1   1
    /     \
   2       6
  / \     / 
 3   1   5   

Example 2:

Input: 
A binary tree as following:
      4
     /   
    2    
   / \   
  3   1    

v = 1

d = 3

Output: 
      4
     /   
    2
   / \    
  1   1
 /     \  
3       1

Note:
The given d is in range [1, maximum depth of the given tree + 1].
The given binary tree has at least one tree node.
输入是一个二叉树，要求在指定的行插入新的一行后返回新的二叉树。

解题思路

典型的二叉树的层序遍历，遍历到指定的深度的时候，返回该层的所有node,执行插入操作。需要注意的是，题目认为树根的level 是1，需要对要求在level 1 插入的情况进行特殊处理。在使用queue 进行树的层序遍历时，每次while循环都记录queue的大小,当前queue存储当前层的所有节点。

C++ 实现

class Solution {
public:
    TreeNode* addOneRow(TreeNode* root, int v, int d) {
        // 处理level 1 的特殊情况
        if (d == 1)
        {
            TreeNode*  new_root = new TreeNode(v);
            new_root->left = root;
            return new_root;
        }
        // 开始层序遍历
        int depth = 2;
        queue<TreeNode*> que;
        que.push(root);
        while (depth < d)
        {
            // 使用queue 的size 记录当前level 的节点个数
            // 每次while循环处理一层
            int size = que.size();
            for (int i = 0; i < size; i++)
            {
                const TreeNode* front = que.front();
                que.pop();
                if (front->left)
                {
                    que.push(front->left);
                }
                if (front->right)
                {
                    que.push(front->right);
                }
            }
            // 处理完一层，计数++
            depth++;
        }
        // 退出循环后 queue中存储的插入的节点
        // 遍历queue 中所有节点，执行插入
        while (!que.empty())
        {
            TreeNode* const node = que.front();
            que.pop();
            TreeNode* left = node->left;
            TreeNode* right = node->right;

            TreeNode* new_node = new TreeNode(v);
            node->left = new_node;
            new_node->left = left;

            new_node = new TreeNode(v);
            node->right = new_node;
            new_node->right = right;

        }

        return root;
    }
};