Problem statement

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

Solution

I saw several really nice and short solutions for this problem, howerer sometimes it is a quite difficult to digest them. Here I tried to find the solution, which can be easily reproduced during real-time interview. The idea as in other solutions is to use recursion.

Let use helper(node) function which return the first and the last elements of linked lists, constucted from subtree of node. Then we can have several cases:

  1. If there is no left and no right children, then we just return (node, node), because in this case the list for node has only one element. node ->
  2. If there is no left children but there is right children, we have a case node -> b2 -> ... -> e2 ->. Then we need to create new connection node -> b2.
  3. If there is left children and there is no right children, we have a case node -> b1 -> ... -> e1 ->, and we need to create connection node -> b1.
  4. If there is left and right children, we have a case node -> b1 -> ... -> e1 -> b2 -> ... -> e2 ->. Then we need to create two more connections: node -> b1 and e1 -> b2.

Complexity

Time complexity is O(n) to traverse our tree once. Space complexity we can say O(h), because we reuse existing nodes with recursion stack.

Code

class Solution:
    def flatten(self, root):
        def helper(node):
            if not node: return (None, None)
            b1, e1 = helper(node.left)
            b2, e2 = helper(node.right)
            node.left = None
            
            if not e1 and not e2:
                return (node, node)
            if not e1 and e2: 
                node.right = b2
                return (node, e2)
            if e1 and not e2:
                node.right = b1
                return (node, e1)
            else:
                node.right = b1
                e1.right = b2
                return (node, e2)

        helper(root)