【Leetcode_总结】623. 在二叉树中增加一行 - python

Q:

给定一个二叉树，根节点为第1层，深度为 1。在其第 d 层追加一行值为 v 的节点。

添加规则：给定一个深度值 d （正整数），针对深度为 d-1 层的每一非空节点 N，为 N 创建两个值为 v 的左子树和右子树。

将 N 原先的左子树，连接为新节点 v 的左子树；将 N 原先的右子树，连接为新节点 v 的右子树。

如果 d 的值为 1，深度 d - 1 不存在，则创建一个新的根节点 v，原先的整棵树将作为 v 的左子树。

示例 1:

输入: 
二叉树如下所示:
       4
     /   \
    2     6
   / \   / 
  3   1 5   

v = 1

d = 2

输出: 
       4
      / \
     1   1
    /     \
   2       6
  / \     / 
 3   1   5   

链接：https://leetcode-cn.com/problems/add-one-row-to-tree/description/

思路：利用层次遍历的思想，遍历到d-1层，然后构造新节点，并插入

代码：

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def addOneRow(self, root, v, d):
        """
        :type root: TreeNode
        :type v: int
        :type d: int
        :rtype: TreeNode
        """
        que = [root]
        level = 1
        if d == 1:
            root_new = TreeNode(v)
            root_new.left = root
            return root_new
        while que:
            if level == d-1:
                for i in range(len(que)):
                    inputs_node_left = TreeNode(v)
                    inputs_node_right = TreeNode(v)
                    node = que.pop(0)
                    temp_letf = None
                    temp_right = None
                    if node.left:
                        temp_letf = node.left
                        node.left = inputs_node_left
                        inputs_node_left.left = temp_letf
                    else:
                        node.left = inputs_node_left
                    if node.right:
                        temp_right = node.right
                        node.right = inputs_node_right
                        inputs_node_right.right = temp_right
                    else:
                         node.right = inputs_node_right             
            else:
                level+=1
                for i in range(len(que)):
                    node = que.pop(0)
                    if node.left:
                        que.append(node.left)
                    if node.right:
                        que.append(node.right)
        return root