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.
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
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)