Flatten Binary Tree to Linked List

flatten_binary_tree.py

"""
Given a binary tree, flatten it to a linked list in-place.

For example, given the following tree:

    1
   / \
  2   5
 / \   \
3   4   6
The flattened tree should look like:

1
 \
  2
   \
    3
     \
      4
       \
        5
         \
          6

Time Complexity: O(N) since we process each node of the tree exactly once.
Space Complexity: O(N) which is occupied by the recursion stack. The problem statement doesn't mention anything about
the tree being balanced or not and hence, the tree could be e.g. left skewed and in that case the longest branch
(and hence the number of nodes in the recursion stack) would be N.
"""


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def flatten(self, root: TreeNode) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        self.helper(root)

    def helper(self, node: TreeNode) -> TreeNode:
        # null edge case
        if not node:
            return node
        # leaf node
        if not node.left and not node.right:
            return node

        # traverse left and right
        left_tail = self.helper(node.left)
        right_tail = self.helper(node.right)

        # if there was a lef subtree, shuffle the connections, so that nothing remains on left
        if left_tail:
            left_tail.right = node.right
            node.right = node.left
            node.left = None

        # return rightmost node once all got connected
        return right_tail if right_tail else left_tail