【LeetCode OJ】Flatten Binary Tree to Linked List

Problem Link:

http://oj.leetcode.com/problems/flatten-binary-tree-to-linked-list/

The problem is asking for flatterning a binary tree to linked list by the pre-order, therefore we could flatten tree from the root. For each node, we link it with its next child in the pre-order, and leave the other child for the furture flattening. We could do this by using a stack. The python code is as follows.

class Solution:
    # @param root, a tree node
    # @return nothing, do it in place
    def flatten(self, root):
        """
        Traverse the tree in pre-order,
        for each node, we link it and its previous node.
        """
        # Special case
        if not root:
            return
        # Previous node
        prev = root
        # Unlinked node in stack
        stack = []
        # Check root's children
        if root.right:
            stack.append(root.right)
        if root.left:
            stack.append(root.left)
        # Link all nodes in the tree
        while stack:
            # Get the next pre-order node
            next_right = stack.pop()
            # Flatten the previous node
            prev.right = next_right
            prev.left = None
            # Go to next pre_order node, and flatten its sub-tree
            prev = next_right
            # Check its sub-tree
            if prev.right:
                stack.append(prev.right)
            if prev.left:
                stack.append(prev.left)