A solution with Python generators ;-)


  • 2
    R

    I split the problem into two parts:

    • A generator that create a sequence of nodes from a pre-order traversal.
    • The flatten method uses the generator to fix the tree, just link prev.right with curr.

    Is recursive, but nice, and the code is reusable, easy to understand and maintain.

    Complessity: O(n)?

    class Solution:
    
        # @param root, a tree node
        # @return nothing, do it in place
        def flatten(self, root):
            
            if not root: return
            
            prev = None
            for node in self.preOrder(root):
                
                node.left = None
                if prev: prev.right = node
                prev = node
        
        
        def preOrder(self, node):
            
            left = node.left
            right = node.right
            
            yield node
            if left:
                for node in self.preOrder(left):
                    yield node
            
            if right:
                for node in self.preOrder(right):
                    yield node

  • 4
    U

    Thanks, your solution let me, another Python beginner, know the existence of generator. It's a nice feature that avoids extra space otherwise needed to contain all the results. For this problem though, it still can be solved in old way.

    class Solution:
    # @param root, a tree node
    # @return nothing, do it in place
    def flatten(self, root):
        if not root:
            return
        pre = None
        stack = [root]
        while stack:
            r = stack.pop()
            if pre:
                pre.left = None
                pre.right = r
            pre = r
            if r.right:
                stack.append(r.right)
            if r.left:
                stack.append(r.left)

  • 0
    S

    one small trick: if you initialize pre = TreeNode(0), then you don't need to check if pre: in the loop. ;-)


Log in to reply
 

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