114 二叉樹展開爲鏈表（python）

LeetCode 114 : 二叉樹展開爲鏈表

 

給定一個二叉樹，原地將它展開爲鏈表。

例如，給定二叉樹

    1
   / \
  2   5
 / \   \
3   4   6
將其展開爲：

1
 \
  2
   \
    3
     \
      4
       \
        5
         \
          6

思路：非遞歸先序遍歷，遍歷過程中維護一個父節點（pre），邊遍歷邊原地修改, 每次連接父節點的右孩子, 將父節點的左孩子 --> None

定義二叉樹：

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

由列表生成二叉樹：

def Create_tree(root,list,i):
    if i<len(list):
        if not list[i]: return None
        else:
            root = TreeNode(list[i])
            root.left = Create_tree(root.left, list, i*2+1)
            root.right = Create_tree(root.right, list, i*2+2)
            return root
    return root

實現代碼：
 

class Solution:
    def flatten(self, root: TreeNode) -> None:
        stack = [root]
        pre = None
        while stack:
            node = stack.pop()
            if pre:
                pre.right = node
                pre.left = None
            if node:
                stack.append(node.right)
                stack.append(node.left)
                pre = node
        return root

測試代碼：

if __name__ == "__main__":
    s = Solution()
    nums = [1,2,5,3,4,None,6]
    root = Create_tree(None, nums, 0)

    res = pre_order2(root)
    print(res)

    # r = s.flatten(root)
    # # 輸出鏈表
    # while r:
    #     print(r.val)
    #     r = r.right