1. 程式人生 > >LeetCode 623 在二叉樹中增加一行 Python

LeetCode 623 在二叉樹中增加一行 Python

給定一個二叉樹,根節點為第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

示例 2:

輸入: 
二叉樹如下所示:
      4
     /   
    2    
   / \   
  3   1    

v = 1

d = 3

輸出: 
      4
     /   
    2
   / \    
  1   1
 /     \  
3       1

注意:

  1. 輸入的深度值 d 的範圍是:[1,二叉樹最大深度 + 1]。
  2. 輸入的二叉樹至少有一個節點
    class Solution:
        def addOneRow(self, root, v, d):
            """
            :type root: TreeNode
            :type v: int
            :type d: int
            :rtype: TreeNode
            """
            if d == 1:
                tmp = TreeNode(v)
                tmp.left = root
                return tmp
            
            #討論d>1的情形
            queue = [(root,1)]
            #首先要將所有d-1層的節點都加進來
            while queue:
                item = queue.pop(0)
                (TmpNode,depth) = item
                if depth >= d-1: 
                    queue.insert(0,(TmpNode,depth))
                    break
                if TmpNode.left: queue.append((TmpNode.left,depth+1))
                if TmpNode.right: queue.append((TmpNode.right,depth+1))
            #此時所有depth=d-1的節點都放進了queue
            
            for (node,depth) in queue:
                tmp = node.left
                mynode = TreeNode(v)
                node.left = mynode
                mynode.left = tmp
                
                tmp = node.right
                mynode = TreeNode(v)
                node.right = mynode
                mynode.right = tmp
            
            return root