leetcode-99. 恢复二叉搜索树

题目

二叉搜索树中的两个节点被错误地交换。

请在不改变其结构的情况下，恢复这棵树。

示例 1:

输入: [1,3,null,null,2]

   1
  /
 3
  \
   2

输出: [3,1,null,null,2]

   3
  /
 1
  \
   2

示例 2:

输入: [3,1,4,null,null,2]

  3
 / \
1   4
   /
  2

输出: [2,1,4,null,null,3]

  2
 / \
1   4
   /
  3

解题思路

$o(n)$的空间复杂度……不想改进了。

两遍中序遍历。
第一遍中序遍历，得到乱序的val，对数组进行排序。
第二遍中序遍历，按照排序后的数组，对树节点赋值。

代码

# 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 recoverTree(self, root: TreeNode) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        stack = [(root, 0)]
        node_val = []
        # inorder
        while stack:
            p, stat = stack.pop()
            if stat == 0:
                if p.right:
                    stack.append((p.right, 0))
                stack.append((p, 1))
                if p.left:
                    stack.append((p.left, 0))
            else:
                node_val.append(p.val)
        # inorder
        node_val.sort()
        stack = [(root, 0)]
        index = 0
        while stack:
            p, stat = stack.pop()
            if stat == 0:
                if p.right:
                    stack.append((p.right, 0))
                stack.append((p, 1))
                if p.left:
                    stack.append((p.left, 0))
            else:
                p.val = node_val[index]
                index += 1