算法题--修正只有2个元素错置的二叉搜索树

image.png

0. 链接

1. 题目

Two elements of a binary search tree (BST) are swapped by mistake.

Recover the tree without changing its structure.

Example 1:

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

   1
  /
 3
  \
   2

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

   3
  /
 1
  \
   2

Example 2:

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

  3
 / \
1   4
   /
  2

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

  2
 / \
1   4
   /
  3

Follow up:

A solution using O(n) space is pretty straight forward.
Could you devise a constant space solution?

2. 思路1: Morris中序遍历

建立螺纹二叉树(Threaded Binary Tree), 即满足下列两个要求：

1.1 所有原本为空的右子节点指向中序遍历顺序之后的那个节点;
1.2 所有原本为空的左子节点指向中序遍历顺序之前的那个节点

具体实现过程:

(1) 初始化指针cur指向root
(2) 当cur不为空时
    -- 如果cur没有左子节点, 说明找到了目前未遍历到的部分的最小值，则
        a) 得到一个中序遍历的元素cur.val
        b) 将cur指向其右子节点(可理解为目前未遍历到的最靠左的那个树的右子节点, 变成了最新的最左端)
    -- 否则
        将pre指针指向cur的左子树中的最右子节点, 不能指向cur(目的是实现1.1的要求)
        ** 若pre的右子节点为空, 说明找到了1.1要求的那个节点, 则
            a) 将pre的右子节点指向cur
            b) 将cur指向cur.left
        ** 否则
            a) 将pre的右子节点置空
            b) 得到一个中序遍历的元素cur.val
            c) 将cur指向其右子节点

对于合法的二叉搜索树而言,

进行中序遍历后得到的应该是一个升序排列的数组, 利用这个特性，在遍历的过程中，用比较前一个得到的值跟当前值，若last_element.val >= cur.val, 则说明至少发现了一个嫌疑点, 用first_element记录此时的last_element, 用second_element记录此时的cur节点；若后续又发现了这种情况的存在, 则更新second_element
交换first_element和second_element的val即可.

3. 代码

# coding:utf8


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


class Solution:
    def recoverTree(self, root: TreeNode) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        first_element = second_element = None
        cur = root
        last_node = None
        while cur is not None:
            if cur.left is None:
                if last_node is not None and cur.val <= last_node.val:
                    if first_element is None:
                        first_element = last_node
                        second_element = cur
                    else:
                        second_element = cur
                last_node = cur
                cur = cur.right
            else:
                pre = cur.left
                while pre.right is not None and pre.right != cur:
                    pre = pre.right
                if pre.right is None:
                    pre.right = cur
                    cur = cur.left
                else:
                    pre.right = None
                    if last_node is not None and cur.val <= last_node.val:
                        if first_element is None:
                            first_element = last_node
                            second_element = cur
                        else:
                            second_element = cur
                    last_node = cur
                    cur = cur.right

        temp = first_element.val
        first_element.val = second_element.val
        second_element.val = temp


def pre_order_traverse(root, array):
    if root is None:
        array.append(None)
        return
    array.append(root.val)
    pre_order_traverse(root.left, array)
    pre_order_traverse(root.right, array)


def print_tree(root):
    results = list()
    pre_order_traverse(root, results)
    print(results)


root1 = node = TreeNode(1)
node.left = TreeNode(3)
node.left.right = TreeNode(2)

root2 = node = TreeNode(4)
node.left = TreeNode(6)
node.right = TreeNode(2)
node.left.right = TreeNode(3)
node.right.left = TreeNode(5)

root3 = node = TreeNode(3)
node.left = TreeNode(1)
node.right = TreeNode(4)
node.right.left = TreeNode(2)

solution = Solution()

print_tree(root1)
solution.recoverTree(root1)
print_tree(root1)
print('=' * 50)
print('')

print_tree(root2)
solution.recoverTree(root2)
print_tree(root2)
print('=' * 50)
print('')

print_tree(root3)
solution.recoverTree(root3)
print_tree(root3)
print('=' * 50)
print('')

输出结果

[1, 3, None, 2, None, None, None]
[3, 1, None, 2, None, None, None]
==================================================

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

[3, 1, None, None, 4, 2, None, None, None]
[2, 1, None, None, 4, 3, None, None, None]
==================================================

结果