二叉树展开为列表（LeetCode）

题目

给你二叉树的根结点 root ，请你将它展开为一个单链表：

• 展开后的单链表应该同样使用 TreeNode ，其中 right 子指针指向链表中下一个结点，而左子指针始终为 null 。
• 展开后的单链表应该与二叉树 先序遍历 顺序相同。

解题

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


def build_tree(nodes, index=0):
    if index >= len(nodes) or nodes[index] is None:
        return None
    root = TreeNode(nodes[index])
    root.left = build_tree(nodes, 2 * index + 1)
    root.right = build_tree(nodes, 2 * index + 2)
    return root


def flatten(root):
    # 使用前序遍历展开二叉树
    if not root:
        return

    stack = [root]
    prev = None

    while stack:
        curr = stack.pop()

        if prev:
            prev.right = curr
            prev.left = None

        if curr.right:
            stack.append(curr.right)
        if curr.left:
            stack.append(curr.left)

        prev = curr

    # 转换为符合题目要求的列表形式
    flattened_list = []
    while root:
        flattened_list.append(root.val)
        root = root.right
        if root:
            flattened_list.append(None)

    return flattened_list


# 测试
nodes = [1, 2, 5, 3, 4, None, 6]
root = build_tree(nodes)
flattened_list = flatten(root)
print(flattened_list)  # 输出：[1, None, 2, None, 3, None, 4, None, 5, None, 6]