Elegant python code using "yield" iterator


  • 3
    M

    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
                    else:
                        second = node
                        break
                pred = node
    
            first.val, second.val = second.val, first.val
    
            return root

  • 0
    C
    This post is deleted!

  • 0
    W

    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.