【劍指 の 精選】詳解「二叉樹中序遍歷的下一個結點」兩種解法 ｜Python 主題月

本文正在參加「Python主題月」，詳情查看 活動連接

題目描述

這是「牛客網」上的「JZ 57 二叉樹的下一個結點」，難度爲「中等」。

Tag : 「劍指 Offer」、「二叉樹」、「中序遍歷」

給定一個二叉樹其中的一個結點，請找出中序遍歷順序的下一個結點而且返回。

注意，樹中的結點不只包含左右子結點，同時包含指向父結點的 next 指針。

下圖爲一棵有 個節點的二叉樹。樹中從父節點指向子節點的指針用實線表示，從子節點指向父節點的用虛線表示。

示例:

輸入:{8,6,10,5,7,9,11},8

返回:9

解析:這個組裝傳入的子樹根節點，其實就是整顆樹，中序遍歷{5,6,7,8,9,10,11}，根節點8的下一個節點就是9，應該返回{9,10,11}，後臺只打印子樹的下一個節點，因此只會打印9，以下圖，其實都有指向左右孩子的指針，還有指向父節點的指針，下圖沒有畫出來。

複製代碼

輸入描述：輸入分爲2段，第一段是總體的二叉樹，第二段是給定二叉樹節點的值，後臺會將這2個參數組裝爲一個二叉樹局部的子樹傳入到函數GetNext裏面，用戶獲得的輸入只有一個子樹根節點。

返回值描述：返回傳入的子樹根節點的下一個節點，後臺會打印輸出這個節點。

示例1

輸入：{8,6,10,5,7,9,11},8

返回值：9

複製代碼

示例2

輸入：{8,6,10,5,7,9,11},6

返回值：7

複製代碼

示例3

輸入：{1,2,#,#,3,#,4},4

返回值：1

複製代碼

示例4

輸入：{5},5

返回值："null"

說明：不存在，後臺打印"null" 

複製代碼

要求：

• 時間：1 s

• 空間：64 M

樸素解法

一個樸素的作法是，根據題目對於 TreeLinkNode 的定義，利用 next 屬性存儲「當前節點的父節點」這一特色。

從入參節點 pNode 出發，不斷利用 next 屬性往上查找，直到找到整棵樹的頭節點，令頭節點爲 root。

而後實現二叉樹的「中序遍歷」，將遍歷過程當中訪問的節點存放到列表 list 中，以後再對 list 進行遍歷，找到 pNode 所在的位置 idx，便可肯定 pNode 是否存在「下一個節點」以及「下一節點」是哪個。

Java 代碼：

import java.util.*;
public class Solution {
    List<TreeLinkNode> list = new ArrayList<>();
    public TreeLinkNode GetNext(TreeLinkNode pNode) {
        // 根據傳入的節點的 next 指針一直往上找，直到找到根節點 root
        TreeLinkNode root = pNode;
        while (root.next != null) root = root.next;

        // 對樹進行一遍「中序遍歷」，保存結果到 list 中
        dfs(root);

        // 從 list 中找傳入節點 pNode 的位置 idx
        int n = list.size(), idx = -1;
        for (int i = 0; i < n; i++) {
            if (list.get(i).equals(pNode)) {
                idx = i;
                break;
            }
        }
        
        // 若是 idx 不爲「中序遍歷」的最後一個元素，那麼說明存在下一個節點，從 list 查找並返回
        // 這裏的 idx == -1 的判斷屬於防護性編程
        return idx == -1 || idx == n - 1 ? null : list.get(idx + 1);
    }
    void dfs(TreeLinkNode root) {
        if (root == null) return;
        dfs(root.left);
        list.add(root);
        dfs(root.right);
    }
}
複製代碼

Python 3 代碼：

class Solution:
    lt = []
    def GetNext(self, pNode):
        # 根據傳入的節點的 next 指針一直往上找，直到找到根節點 root
        root = pNode
        while root.next is not None:
            root = root.next
        # 對樹進行一遍「中序遍歷」，保存結果到 list 中
        self.dfs(root)
        
        # 從 list 中找傳入節點 pNode 的位置 idx
        n = len(self.lt)
        idx = -1
        for i in range(n):
            if self.lt[i] == pNode:
                idx = i
                break
        # 若是 idx 不爲「中序遍歷」的最後一個元素，那麼說明存在下一個節點，從 list 查找並返回
        # 這裏的 idx == -1 的判斷屬於防護性編程
        return None if idx == -1 or idx == n -1 else self.lt[idx + 1]
    
    def dfs(self, root):
        if root is None:
            return
        self.dfs(root.left)
        self.lt.append(root)
        self.dfs(root.right)
複製代碼

• 時間複雜度：找頭節點 root 最多訪問的節點數量不會超過樹高 h；進行中序遍歷的複雜度爲 O(n)；從中序遍歷結果中找到 pNode 的下一節點的複雜度爲 O(n)。總體複雜度爲 O(n)
• 空間複雜度：忽略遞歸帶來的額外空間消耗。複雜度爲 O(n)

進階解法

另一個「進階」的作法是充分利用「二叉樹的中序遍歷」來實現的。

咱們知道，二叉樹「中序遍歷」的遍歷順序爲 「左-根-右」。

能夠根據傳入節點 pNode 是否有「右兒子」，以及傳入節點 pNode 是否爲其「父節點」的「左兒子」來進行分狀況討論：

• 傳入節點 pNode 有「右兒子」：根據「中序遍歷」的遍歷順序爲 「左-根-右」，能夠肯定「下一個節點」必然爲當前節點的「右子樹」中「最靠左的節點」；

• 傳入節點 pNode 沒有「右兒子」，這時候須要根據當前節點是否爲其「父節點」的「左兒子」來進行分狀況討論：

• 若是傳入節點 pNode 自己是其「父節點」的「左兒子」，那麼根據「中序遍歷」的遍歷順序爲爲 「左-根-右」 可知，下一個節點正是該父節點，直接返回該節點便可；

• 若是傳入節點 pNode 自己是其「父節點」的「右兒子」，那麼根據「中序遍歷」的遍歷順序爲爲 「左-根-右」 可知，其父節點已經被遍歷過了，咱們須要遞歸找到符合 node.equals(node.next.left) 的節點做爲答案返回，若是沒有則說明當前節點是整顆二叉樹最靠右的節點，這時候返回 null 便可。

代碼：

public class Solution {
    public TreeLinkNode GetNext(TreeLinkNode pNode) {
        if (pNode.right != null) {
            // 若是當前節點有右兒子，下一節點爲當前節點「右子樹中最靠左的節點」
            pNode = pNode.right;
            while (pNode.left != null) pNode = pNode.left;
            return pNode;
        } else {
            // 若是當前節點沒有右兒子，則「往上找父節點」，直到出現知足「其左兒子是當前節點」的父節點
            while (pNode.next != null) {
                if (pNode.equals(pNode.next.left)) return pNode.next;
                pNode = pNode.next;
            }
        }
        return null;
    }
}
複製代碼

Python 3 代碼：

class Solution:
    def GetNext(self, pNode):
        if pNode.right is not None:
            # 若是當前節點有右兒子，下一節點爲當前節點「右子樹中最靠左的節點」
            pNode = pNode.right
            while pNode.left is not None:
                pNode = pNode.left
            return pNode
        else:
            # 若是當前節點沒有右兒子，則「往上找父節點」，直到出現知足「其左兒子是當前節點」的父節點
            while pNode.next is not None:
                if pNode == pNode.next.left:
                    return pNode.next
                pNode = pNode.next
        return None
複製代碼

• 時間複雜度：O(n)
• 空間複雜度：O(1)

最後

這是咱們「劍指 の 精選」系列文章的第 No.57 篇，系列開始於 2021/07/01。

該系列會將牛客網「劍指 Offer」中比較經典而又不過期的題目都講一遍。

在提供追求「證實」&「思路」的同時，提供最爲簡潔的代碼。

歡迎關注，交個朋友 (｀・ω・´)