LeetCode 99 | 如何不用递归遍历二叉搜索树？MT方法给你答案

今天是LeetCode专题的第64篇文章，我们来看LeetCode第99题BST（二叉搜索树）的还原问题。

这题的官方难度是Hard，点赞1666，反对77，这个反对比例小于5%，称得上是好评如潮。通过率是39.6%，从这个通过率来看，其实在Hard难度的问题当中不算很高。实际上也的确如此，这道题同样非常看重基础，考察的是我们对数据结构的理解以及代码的编写能力。

题意

有一颗二叉搜索树被交换了其中的两个元素，导致破坏了二叉搜索树的性质。现在给定我们交换之后的二叉搜索树，要求我们还原它得到交换之前的结果。

样例

Input: [1,3,null,null,2]

   1
  /
 3
  \
   2

Output: [3,1,null,null,2]

   3
  /
 1
  \
   2

第二个样例

Input: [3,1,4,null,null,2]

  3
 / \
1   4
   /
  2

Output: [2,1,4,null,null,3]

  2
 / \
1   4
   /
  3

提示

• 算法复杂度应该是 $O(n)$
• 能否让空间复杂度维持在常量级别( $O(1)$ )

题解

题意已经说得很清楚了，在整个二叉树当中一共只有两个元素交换了位置导致了错误。那么我们要将这棵二叉树还原，需要首先找到这两个交换了位置的元素，找到了元素之后就方便了，只需要交换它们就可以了。

但问题来了，我们判断BST是否合法容易，但是我们怎么寻找摆放错误的元素呢？

比如说我们知道了以u为根节点的BST是非法的，非法的原因是因为u的值大于右子树中的最小值。其实这时候有两种可能，一种是右子树的最小值摆放错误了，还有一种可能是u本身摆放错了。虽然我们知道只有两个元素摆放错了，但是要通过递归将它们找出来却不太容易，似乎不能直接得到答案。

关于这里的思路我也思考了很久，直到找到了一个点解开了这一切。这个破题的点在哪里呢？在中序遍历。

对于一棵合法的BST它中序遍历的结果应该是升序的，想到这里剩下的就迎刃而解了。因为我们已经知道了只有两个元素的位置错了，那么我们只需要找到一个序列当中不满足升序的两个元素，它们必然就是摆放错误的点。最后只需要交换它们的值就行了。

如果你想不到中序遍历，那么这道题会非常困难，尤其是如果你尝试用递归去解的话，你会发现这个递归代码怎么也写不对。这个时候尤其需要头脑冷静，目标越接近越要好好思考，是不是真的只差临门一脚了，否则一位坚持，也只会浪费时间。

想明白了这点之后，还剩下最后一个问题，就是我们怎么在一个交换了两个元素的升序序列当中找到这两个元素。我们需要分情况讨论，什么情况呢，就是这两个摆放错误的元素是否相邻。

如果这两个元素相邻，那么我们只会找到一处顺序不对的地方。举个例子[1, 3, 2, 4, 5]，这里发生错位的是2和3，我们寻找所有a[i] < a[i-1]的i只能找到一个。

如果这两个元素不相邻，那么我们在遍历的时候可以找到两个错位的地方。比如[1, 4, 3, 2, 5]，我们会发现3和2都小于它前位。但是实际上错位的是2和4，而不是2和3。

也就是说如果我们只遇到一次顺序颠倒的情况，那么颠倒的两个数位置都错了。如果遇到了不止两次，那么就是第一次的前位和第二次的后位错了。把这个情况理清楚了，代码其实很简单。

# 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.
        """
        inorder = []

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

            dfs(node.left)
            inorder.append(node)
            dfs(node.right)

        # 中序遍历获得所有元素
        dfs(root)
        n = len(inorder)
        x, y = None, None
        for i in range(n-1):
            if inorder[i+1].val < inorder[i].val:
                # 我们用y记录后一个错误的位置
                # 如果遇到两次，那么取最后出现的y
                y = inorder[i+1]
                # x用来记录第一次错位时的前位元素
                # 如果遇到两次错位，直接break
                if x is None:
                    x = inorder[i]
                else:
                    break


        x.val, y.val = y.val, x.val

进阶解法

上面的解法虽然能AC但是并不完全满足题目的要求，不满足的地方在于题目当中要求我们的空间复杂度是 $O(1)$ ，但是我们用数组存储了中序遍历的结果。

但是我们不记录下来又怎么样才能知道哪里的元素出现了乱序呢？

要解决这个问题，需要用到一种特殊的遍历二叉树的算法，称为Morris Traversal方法。算法原理可以参考这篇博文：https://blog.csdn.net/u013007900/article/details/77663733

简单来说就是我们在遍历二叉树的时候先派遣一个指针pnt，它用来遍历左子树的最右侧树枝上的根节点，也就是找到左子树当中最大元素的位置。然后从这个位置建立一个指针指向当前的根节点。这样当我们遍历到pnt的位置的时候就可以通过pnt新建的指针回到根节点。这样我们就可以用迭代的方式以中序遍历的顺序遍历整棵二叉搜索树。

具体的过程可以参考一下下图，应该非常清晰：

对于当前节点cur来说，先遣指针pnt指向的节点其实就是中序遍历之后在它前一位相邻的节点。所以和上面的逻辑一样，我们只需要比较一下pnt是否大于cur就可以知道是否有错位的情况出现。

这两个算法的内核逻辑是完全一样的，唯一不同的是上面的方法是我们先遍历获得完所有的数据之后再来寻找错位的点。而这种算法是一边遍历一边寻找。

Morris Traversal方法只有一个相对比较难理解的点，就是当我们通过pnt找到左子树最大值的时候，我们需要判断一下pnt是否已经连过了边到cur。这么做的原因是我们无法判断cur是第一次遇见，还是之前我们已经执行过了一次遍历，又通过pnt回到了cur。

我这么说有一点抽象， 我们可以来看下上图当中的这个部分：

这个部分话的逻辑是我们已经遍历完了左子树当中的所有元素，接着我们会通过cur的右指针回到6的位置。于是接下来我们会遍历6，但一个问题是我们此时无法判断6的左子树是否已经遍历过了。唯一的判断方法就是通过5这个位置存不存在指向6的指针来判断。如果存在，那么说明之前已经遍历过了，我们需要断开这个指针，并且开始遍历6的右子树。

如果你能理解这一点的话，那么这个算法应该已经没多大问题了。

我们来看代码：

# 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.
        """

        x, y, pnt = None, None, None
        pre = None

        while root is not None:
            # 如果存在左子树那么找到pnt
            if root.left is not None:
                pnt = root.left
                while pnt.right is not None and pnt.right.val != root.val:
                    pnt = pnt.right

                # 如果pnt的right是空，说明root是第一次遍历到
                if pnt.right is None:
                    pnt.right = root
                    root = root.left
                else:
                    # 否则比较前驱和root的大小
                    # 这里要注意没有break，因为我们遍历完之后需要去掉多余的指针，否则会导致死循环
                    if pre is not None and pre.val > root.val:
                        y = root
                        if x is None:
                            x = pre
                    # 将root更新成前驱
                    pre = root
                    pnt.right = None
                    root = root.right
            else:
                # 向右移动比较简单，只需要考虑前驱的情况
                if pre is not None and pre.val > root.val:
                    y = root
                    if x is None:
                        x = pre
                pre = root
                root = root.right

        x.val, y.val = y.val, x.val

今天的题目相比之前难度稍稍大了一些，尤其是在最后一种解法很难想到，即使看了题解理解起来也不容易。这也是非常正常的，如果觉得理解吃力，建议可以回过头多研究一下上面的遍历图，说不定会有豁然开朗的感觉。

今天的文章到这里就结束了，如果喜欢本文的话，请来一波素质三连，给我一点支持吧（关注、在看、点赞）。