LeetCode #99: Recover Binary Search Tree

Problem Statement

(Source) Two elements of a binary search tree (BST) are swapped by mistake. Recover the tree without changing its structure.

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

Analysis

Think of the in-order traversal sequence of a binary search tree. It is an increasing sequence. If two elements of the original binary search tree are swapped, then the new in-order traversal sequence would not be in ascending order anymore, but it must fit into one of the following two cases:
(1) There exists only one pair of elements (seq[i], seq[i + 1]) in the new sequence such that seq[i] > seq[i + 1].
(2) There exists two pairs of elements (seq[i], seq[i + 1]), (seq[j], seq[j + 1]) in the new sequence such that seq[i] > seq[i + 1] && seq[j] > seq[j + 1].

Therefore, the idea to solve the problem is to find the first node corresponding to the first element of the first pair, and the last node corresponding to the last element of the last pair, no matter which case the new sequence fits into, and swap them.

Solution below uses an recursive implementation below, which runs in:
- Time complexity: O(n), where n is the total number of the binary search tree.
- Space complexity: O(1).

Solution

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

class Solution(object):
    def recoverTree(self, root):
        """
        :type root: TreeNode
        :rtype: void Do not return anything, modify root in-place instead.
        """
        swapped_nodes = []
        prev = [None]

        def inorder(p, prev, swapped_nodes):
            if p.left:
                inorder(p.left, prev, swapped_nodes)

            if prev[0] and prev[0].val > p.val:
                if len(swapped_nodes) == 0:
                    swapped_nodes.extend([prev[0], p])
                else:
                    swapped_nodes[1] = p
            prev[0] = p

            if p.right:
                inorder(p.right, prev, swapped_nodes)

        inorder(root, prev, swapped_nodes)

        swapped_nodes[0].val, swapped_nodes[1].val = swapped_nodes[1].val, swapped_nodes[0].val