Leetcode (Python): Recover Binary Search Tree

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?

Solution

If we traverse the tree inorder, we just need to identify the 2 elements that are misplaced and swap their values. We also keep a third pointer that points to the previous element as it is needed to identify when an element is misplaced.

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

class Solution:
    # @param root, a tree node
    # @return a tree node
    
    def recoverTree(self, root):
        self.pre = None
        self.node1 = None
        self.node2 = None
        self.inOrder(root)
        val = self.node1.val
        self.node1.val = self.node2.val
        self.node2.val = val
        return root
        
    def inOrder(self, root):
        if root == None:
            return
        self.inOrder(root.left)
        if self.pre == None:
            self.pre = root
        if self.node1 == None and self.pre.val > root.val:
            self.node1 = self.pre
            self.node2 = root
        elif self.pre.val > root.val:
            self.node2 = root
        self.pre = root;
        self.inOrder(root.right)
        return