Elegant python code using "yield" iterator

  • 3

    Basic algorithm is inorder traveling the tree to find the swapped pair nodes.

    By using "yield" iterator, the traveling procedure can be very elegant.

    def inorderIter(root):
        if root:
            for node in inorderIter(root.left): 
                yield node
            yield root        
            for node in inorderIter(root.right): 
                yield node
    class Solution:
        # @param root, a tree node
        # @return a tree node
        def recoverTree(self, root):
            pred, first = None, None
            for node in inorderIter(root):
                if pred and pred.val > node.val:
                    if first == None:
                        first = pred
                        second = node
                        second = node
                pred = node
            first.val, second.val = second.val, first.val
            return root

  • 0
    This post is deleted!

  • 0

    It seems to be incorrect.

    Input: {0,1}
    Output: Do not return anything, modify the binary tree in-place instead.
    Expected: {1,0}

  • 0

    Such an elegant solution!

    @Weissach need to comment out the last line of code

  • 0

    I am not sure about the memory usage of the iterator. It seems that the iterator uses recursion and is not O(n) space.

Log in to reply

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