算法题--将一棵二叉树变形为只包含右子树

image.png

0. 链接

题目链接

1. 题目

Given a binary tree, flatten it to a linked list in-place.

For example, given the following tree:

    1
   / \
  2   5
 / \   \
3   4   6
The flattened tree should look like:

1
 \
  2
   \
    3
     \
      4
       \
        5
         \
          6

2. 思路1: 反向后序遍历+递归

• 就是按照反向后序遍历的方法, 左子树变形后的结果作为root的右节点, 而root原来的右子树, 接在后面

3. 代码

# coding:utf8


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def flatten(self, root: TreeNode) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        pre = None

        def update(node):
            if node is None:
                return

            nonlocal pre

            update(node.right)
            update(node.left)
            node.right = pre
            node.left = None
            pre = node

        update(root)


def print_tree(root):
    def traverse(node, array):
        array.append(node.val)
        if node.left is not None:
            traverse(node.left, array)
        else:
            array.append(None)
        if node.right is not None:
            traverse(node.right, array)
        else:
            array.append(None)

    array = list()
    if root is not None:
        traverse(root, array)
    print(array)


solution = Solution()

root1 = node = TreeNode(1)
node.left = TreeNode(2)
node.right = TreeNode(5)
node.left.left = TreeNode(3)
node.left.right = TreeNode(4)
node.right.right = TreeNode(6)
print_tree(root1)
solution.flatten(root1)
print_tree(root1)

root2 = node = TreeNode(1)
node.right = TreeNode(2)
node.right.left = TreeNode(3)
print_tree(root2)
solution.flatten(root2)
print_tree(root2)

输出结果

[1, 2, 3, None, None, 4, None, None, 5, None, 6, None, None]
[1, None, 2, None, 3, None, 4, None, 5, None, 6, None, None]
[1, None, 2, 3, None, None, None]
[1, None, 2, None, 3, None, None]

4. 结果