leetcode: 99. Recover Binary Search Tree

Problem

# 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?

AC

class TreeNode():
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution():
    def recoverTree(self, root):
        self.lefty = None
        self.righty = None
        self._minnode = None
        self._maxnode = None
        self.preorder_visit(root)
        self.postorder_visit(root)
        # print self.lefty, self.righty
        self.lefty.val, self.righty.val = self.righty.val, self.lefty.val

    def preorder_visit(self, node):
        # first check left and right trees
        if node.left:
            self.preorder_visit(node.left)
            if self.lefty:
                return
        if self._maxnode and node.val<self._maxnode.val:
            self.lefty = self._maxnode
            return
        if not self._maxnode or node.val > self._maxnode.val:
            self._maxnode = node
        if node.right:
            self.preorder_visit(node.right)

    def postorder_visit(self, node):
        # first check left and right trees
        if node.right:
            self.postorder_visit(node.right)
            if self.righty:
                return
        if self._minnode and node.val > self._minnode.val:
            self.righty = self._minnode
            return
        if not self._minnode or node.val < self._minnode.val:
            self._minnode = node
        if node.left:
            self.postorder_visit(node.left)


# Time:  O(n)
# Space: O(1)
class TreeNode():
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution2():
    def recoverTree(self, root):
        return self.MorrisTraversal(root)

    def MorrisTraversal(self, root):
        if root is None:
            return
        broken = [None, None]
        pre, cur = None, root

        while cur:
            if cur.left is None:
                self.detectBroken(broken, pre, cur)
                pre = cur
                cur = cur.right
            else:
                node = cur.left
                while node.right and node.right != cur:
                    node = node.right

                if node.right is None:
                    node.right =cur
                    cur = cur.left
                else:
                    self.detectBroken(broken, pre, cur)
                    node.right = None
                    pre = cur
                    cur = cur.right

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

        return root

    def detectBroken(self, broken, pre, cur):
        if pre and pre.val > cur.val:
            if broken[0] is None:
                broken[0] = pre
            broken[1] = cur


if __name__ == "__main__":
    root, root.left = TreeNode(0), TreeNode(1)
    assert Solution().recoverTree(root) == None