SIMPLE Python solution, with REAL O(1) space by using Morris traversal (with comments)


  • 1
    def recoverTree(self, root):
        cur, prev, first, second, predec = root, None, None, None, None
        while cur:
            if cur.left is None:  # I have no left child so it's my turn
                if predec is not None and cur.val < predec.val:
                    if not first:
                        first = predec
                    second = cur
                predec = cur
                cur = cur.right  # After visit, go right
            else:
                # I need to find my predecessor and make me its right child
                # Second check in right is to make sure that I don't loop
                prev = cur.left
                while prev.right is not None and prev.right != cur:
                    prev = prev.right
                if prev.right is None:  # Not visited yet
                    prev.right = cur
                    cur = cur.left
                else:  # It's my turn
                    if predec is not None and cur.val < predec.val:
                        if not first:
                            first = predec
                        second = cur
                    predec = cur
                    # I remove the link after having visited
                    prev.right = None
                    cur = cur.right
        first.val, second.val = second.val, first.val

  • 0
    C

    How about using a generator? This is partly the reason why I chose python.

    class Solution(object):
        def recoverTree(self, root):
            prev = cur = None
            first = second = None
            for cur in self.morris_inorder(root):
                if prev and prev.val > cur.val:
                    if not first:
                        first = prev
                    second = cur
                prev = cur
            first.val, second.val = second.val, first.val
    
        def morris_inorder(self, cur):
            while cur:
                if not cur.left:
                    yield cur
                    cur = cur.right
                else:
                    tmp = cur.left
                    while tmp.right and tmp.right != cur:
                        tmp = tmp.right
                    if not tmp.right:
                        tmp.right = cur
                        cur = cur.left
                    else:
                        yield cur
                        cur = cur.right
                        tmp.right = None
    

  • 0
    F

    Any special advantage of using generator here? I think this might be unsafe since the Morris traversal changes the tree structure. If you didn't finish the full traversal (say, got an error within the for loop), the original tree will be left modified.


  • 0
    C

    @fwang As a matter of fact I just turned to python not long ago. In the first place I went for simplicity and readability. With a generator you decouple the logic of tree-traversal and recover the reversed two. One function only does one thing.

    And error-prone, good point. Haven't thought about that before. But in the real world keeping a function nice and simple by making sure it only does one little thing is one good way to prevent bugs. If error is easy to happen in the for-loop already, what would it be like to copy that loop around within the morris-traversal? With a generator you can put that code within a try-finally to make sure you consume the generator. And Thank you for bringing it up, really a good point.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.