Flatten Binary Tree to Linked List

This problem asks us to flatten a binary tree into a linked list in-place, following a pre-order traversal. Each node's right pointer should point to the next node in the pre-order traversal, and the left pointer should be null.

Naive Approach (Using Extra Space)

A straightforward but less efficient approach is to perform a pre-order traversal of the tree, storing the nodes in a list. Then, iterate through the list and update the left and right pointers of each node accordingly.

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right



def flatten_naive(root: TreeNode) -> None:
    """Flattens the tree using extra space to store pre-order traversal.
    "
    if not root:
        return

    nodes = []

    def preorder(node):
        if node:
            nodes.append(node)
            preorder(node.left)
            preorder(node.right)

    preorder(root)

    for i in range(len(nodes) - 1):
        nodes[i].left = None
        nodes[i].right = nodes[i + 1]

    nodes[-1].left = None
    nodes[-1].right = None

Time Complexity: O(N), where N is the number of nodes in the tree (for traversal and pointer updates).

Space Complexity: O(N), to store the nodes in the nodes list.

Optimal Approach (In-Place)

A more efficient in-place approach involves recursion. For each node, we flatten its left and right subtrees. Then, we move the flattened left subtree to the right of the current node and attach the flattened right subtree to the end of the flattened left subtree.

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def flatten(root: TreeNode) -> None:
    """Flattens the tree in-place.
    """
    if not root:
        return

    # Flatten left and right subtrees
    flatten(root.left)
    flatten(root.right)

    # Store the original right subtree
    original_right = root.right

    # Move the flattened left subtree to the right
    if root.left:
        root.right = root.left
        root.left = None

        # Find the tail of the flattened left subtree
        tail = root.right
        while tail.right:
            tail = tail.right

        # Attach the original right subtree to the tail
        tail.right = original_right

Example:

Let's consider the example tree from the prompt:

Flatten left subtree (2 -> 3 -> 4) and right subtree (5 -> 6).
Move flattened left subtree (2 -> 3 -> 4) to the right of 1.
Attach flattened right subtree (5 -> 6) to the end of the flattened left subtree (after 4).

The result is: 1 -> 2 -> 3 -> 4 -> 5 -> 6

Big(O) Runtime Analysis

The flatten function has a time complexity of O(N), where N is the number of nodes in the tree. Although it appears to have nested loops due to the recursion and the while loop for finding the tail, each node is visited and processed a constant number of times. The recursion visits each node once, and the while loop traverses the flattened left subtree, which in total, across all calls, visits each node a constant number of times.

Big(O) Space Usage Analysis

The space complexity of the flatten function is O(H), where H is the height of the tree. This is due to the recursive call stack. In the worst case (a skewed tree), H = N, so the space complexity would be O(N). In the best case (a balanced tree), H = log N, so the space complexity would be O(log N).

Edge Cases

Empty Tree: If the root is None, the function should return without doing anything. The base case if not root: return handles this case.
Single Node Tree: If the tree contains only one node, the function should also return without doing anything, as a single node is already "flattened".
Skewed Tree (Left or Right): The algorithm should handle skewed trees correctly, as it recursively flattens subtrees.

Code with test cases

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def flatten(root: TreeNode) -> None:
    """Flattens the tree in-place.
    """
    if not root:
        return

    # Flatten left and right subtrees
    flatten(root.left)
    flatten(root.right)

    # Store the original right subtree
    original_right = root.right

    # Move the flattened left subtree to the right
    if root.left:
        root.right = root.left
        root.left = None

        # Find the tail of the flattened left subtree
        tail = root.right
        while tail.right:
            tail = tail.right

        # Attach the original right subtree to the tail
        tail.right = original_right


# Helper function to convert flattened tree to list for testing
def tree_to_list(root: TreeNode) -> list:
    result = []
    while root:
        result.append(root.val)
        root = root.right
    return result


# Test cases

# Test case 1: Example tree from prompt
root1 = TreeNode(1, TreeNode(2, TreeNode(3), TreeNode(4)), TreeNode(5, None, TreeNode(6)))
flatten(root1)
assert tree_to_list(root1) == [1, 2, 3, 4, 5, 6]

# Test case 2: Empty tree
root2 = None
flatten(root2)
assert root2 is None

# Test case 3: Single node tree
root3 = TreeNode(1)
flatten(root3)
assert tree_to_list(root3) == [1]

# Test case 4: Skewed left tree
root4 = TreeNode(1, TreeNode(2, TreeNode(3, TreeNode(4, TreeNode(5)))), None)
flatten(root4)
assert tree_to_list(root4) == [1, 2, 3, 4, 5]

# Test case 5: Skewed right tree
root5 = TreeNode(1, None, TreeNode(2, None, TreeNode(3, None, TreeNode(4, None, TreeNode(5)))))
flatten(root5)
assert tree_to_list(root5) == [1, 2, 3, 4, 5]

print("All test cases passed!")

How do you flatten a binary tree into a linked list in-place, following a pre-order traversal?