1 題目描述


給定一個二元樹，然後每個節點有一個 next 指針，將它指向它右邊的節點。如果沒有右邊的節點，則設為 None。

題目有特殊要求，不能使用額外的空間，也就是說不能使用 Queue 來解這題。

2 解法

我們先走簡單路線，你會發現基本上就是把同一層的節點串起來，然後再串下一層的節點。因此我們可以使用BFS的方式來解這題。講到BFS的演算法，可以參考這篇LeetCode 課前預習 - 掌握 BFS 與 DFS 指南。

2.1 Queue

我們可以知道在Iterative的方式下，先不考慮題目限制，我們可以使用Queue來解這題。

BFS的演算法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def bfs_traversal(root):
    def bfs(nodes):
        # 如果沒有節點，就直接返回
        if not nodes:
            return 

        next_level = []     # 裝下一層的每個節點
        result = []         # 裝目前這層的節點值
        
        for node in nodes:  # 循序這層的每個節點
            result.append(node.val)
            if node.left:   
                next_level.append(node.left)    # 把目前這層的子節點，都放到下一層的next_level
            if node.right:  
                next_level.append(node.right)   # 把目前這層的子節點，都放到下一層的next_level
        
        bfs(next_level)     # 繼續執行下一層的節點
    
    if root:
        bfs([root])
    return [node for level in result for node in level]

從上面得到一個概念，那就是如果我們把 next_level也就是當前的 nodes 串起來，不就達到題目的要求了嗎？我們希望在執行 for loop 每個 node 的時候，可以把 nodes之間串起來。可以先建立一個 dummyhead，然後串 node，這樣進入 for loop 的時候，就可以 dummyhead.next 串起來，傳完後再重新把 dummyhead 斷開就好。
大概是這種感覺:

把同level的nodes串起來

1
2
3
4
5
6
7
8
dummy = Node(-1)
cur = dummy
for node in nodes:
    cur.next = node
    cur = cur.next
    ...
dummy.next = None
...

基本上我們就完成了，以下是完整的程式碼：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def connect(self, root: Node) -> Node:
    def bfs(nodes):
        if not nodes:
            return
        next_level = []
        dummy = Node(-1)
        pre = dummy     # dummy head in each level
        for node in nodes:
            pre.next = node
            pre = node
            if node.left:
                next_level.append(node.left)
            if node.right:
                next_level.append(node.right)
        dummy.next = None   # reset dummy
        bfs(next_level)

    if root:
        bfs([root])

2.2 Space O(1)

但是這樣還不夠，因為題目有要求 O(1)的空間限制，因此我們不能使用 Queue 來解這題。我參考了這篇Leetcode wang的文章，我覺得挺清楚的。通常講到 O(1) 就要聯想到 pointer，那我們需要兩個pointer：

  • 一個pointer紀錄上一層的level的node，我們稱作pre;
  • 另一個pointer紀錄當前層level的node，我們稱作cur

pre跟cur都要往右移動，也就是cur負責串當前層的節點，pre負責移動上一層的node，以方便cur繼續連下一個樹同層的節點。大概是這種感覺：

處理pre的子節點時，cur要負責把小孩子都串起來

當cur的小孩子串完後，pre要往右移動，cur才可以繼續串

從以上規律可以發現：

  • dummy 負責串每層的第一個
  • pre 會一直往右走直到觸底，就換下層的第一個，可以指向dummy.next
  • cur 負責把當前的層串起來

完整程式碼如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def connect(self, root: Node) -> Node:
    pre = root 
    while pre: 
        dummy = Node(-1)
        cur = dummy
        while pre:  # pre 如果觸底了，pre 就換下層的第一個
            if pre.left:
                cur.next = pre.left
                cur = cur.next
            if pre.right:
                cur.next = pre.right
                cur = cur.next
            pre = pre.next  # pre 會一直往右走直到觸底
        pre = dummy.next
        dummy.next = None # reset dummy
    return root

3 總結

這題花了大概30分鐘想出來，這是我第一次寫BFS的題目，因此還是有借住一點idea，搭配BFS的演算法，免強想出一個O(n)空間複雜度的解法。這題的重點在於，如何把同一層的節點串起來，然後再串下一層的節點。這樣的題目，我覺得還是要多練習，才能更加熟練。