DFS應用: 在二元樹插入新的一層 Add one row to Tree_Leetcode #623

Ricardo Gomez Angel on Unsplash

題目敘述

題目會給定一顆二元樹的根結點，
要求我們在指定的層樹d，插入新的一層，節點值為v。

測試範例

Example 1:

Input: root = [4,2,6,3,1,5], val = 1, depth = 2
Output: [4,1,1,2,null,null,6,3,1,5]

Example 2:

Input: root = [4,2,null,3,1], val = 1, depth = 3
Output: [4,2,null,1,1,3,null,null,1]

約束條件

Constraints:

• The number of nodes in the tree is in the range [1, 10^4].

節點總數目介於1 ~ 一萬之間。

• The depth of the tree is in the range [1, 10^4].

深度介於1 ~ 一萬之間。

• -100 <= Node.val <= 100

節點值介於-100 ~ 100 之間。

• -10^5 <= val <= 10^5

新插入的節點值介於 負十萬~正十萬之間。

• 1 <= depth <= the depth of tree + 1

新插入的那一層介於1 ~ 原本的層樹+1

演算法 DFS + 共通模式

題目已經開門見山，要求在指定的層樹d，插入新的一層，節點值為v。

示意圖

共通模式:

有一個很直覺的想法，先找到第d層，接著插入指定的新節點v，串接下一層d+1。

接著下方的節點都往後退一層即可。

也就是說，相當於

當d=n時，相當於往下走一層，接著在d-1層插入新節點v。

初始條件:

當d=1時，新插入的位置會取代原本的樹根，原本舊的樹根當作我的左子樹去串接。

當d=2時，剛好是樹根的下方，左、右子樹分別插入新節點。原本舊的節點通通往後退一層。

程式碼 DFS + 共通模式

class Solution:
    def addOneRow(self, root: TreeNode, v: int, d: int) -> TreeNode:
        
        if not root:
            
            # base case: empty tree
            return None
        
        elif d == 1:
            
            # base case: add one row above original root
            return TreeNode(v, left=root, right=None)
        
        elif d == 2:
            
            # base case: add one row just below original root
            root.left = TreeNode(v, left=root.left, right=None)
            root.right = TreeNode(v, left=None, right=root.right)
            
            return root
            
        else:
            
            # general case: depth >= 3
            # do it in DFS with common pattern
            root.left = self.addOneRow(root.left, v, d-1)
            root.right = self.addOneRow(root.right, v, d-1)
            
            return root  

複雜度分析 DFS + 共通模式

時間複雜度:

從上到下掃描到第d層，最深頂多到第n層，每個節點最多拜訪一次，所需時間為O(n)

空間複雜度:

從上到下掃描到第d層，最深頂多到第n層，每個節點最多拜訪一次，所需空間為O(n)

關鍵知識點 找出共通模式

根據題意，找出插入節點的共通模式。

發現可以用遞回結構描述:

在第d層插入新節點v，等價於 先往下走一層，接著在第(d-1)層插入新節點v

Reference:

[1] Add One Row to Tree - LeetCode