Search code examples
algorithmbinary-treetraversaltree-traversal

Reconstruct binary tree from array


Given an array representing the level-order traversal of a binary tree (not BST), reconstruct the tree.

Example:

Input: [6,x,8,4,x]
output: 6
         \
          8
         / 
        4   
x in the array represents null nodes

Many existing solutions use 2i + 1 and 2i + 2 to access the left and right child, where i is the current node in the array. This solution only works if the input array represents a complete binary tree, which means the input in the example above has to be [6,x,8,x,x,4,x]. I need a solution that works with any input. Thank you!

Edit: Assume the TreeNode class has the following structure:

class TreeNode{
  String value;
  TreeNode left;
  TreeNode right;
}

Also, the output was just an illustration of what the tree should look like. But you can just return a TreeNode which is the root of the tree.


Solution

  • Could you do something like this?

    class TreeNode:
        def __init__(self, x):
            self.val = x
            self.left = None
            self.right = None
    
        def display(self):
            lines, *_ = self._display_aux()
            for line in lines:
                print(line)
    
        def _display_aux(self):
            """Returns list of strings, width, height, and horizontal coordinate of the root."""
            # No child.
            if self.right is None and self.left is None:
                line = '%s' % self.val
                width = len(line)
                height = 1
                middle = width // 2
                return [line], width, height, middle
    
            # Only left child.
            if self.right is None:
                lines, n, p, x = self.left._display_aux()
                s = '%s' % self.val
                u = len(s)
                first_line = (x + 1) * ' ' + (n - x - 1) * '_' + s
                second_line = x * ' ' + '/' + (n - x - 1 + u) * ' '
                shifted_lines = [line + u * ' ' for line in lines]
                return [first_line, second_line] + shifted_lines, n + u, p + 2, n + u // 2
    
            # Only right child.
            if self.left is None:
                lines, n, p, x = self.right._display_aux()
                s = '%s' % self.val
                u = len(s)
                first_line = s + x * '_' + (n - x) * ' '
                second_line = (u + x) * ' ' + '\\' + (n - x - 1) * ' '
                shifted_lines = [u * ' ' + line for line in lines]
                return [first_line, second_line] + shifted_lines, n + u, p + 2, u // 2
    
            # Two children.
            left, n, p, x = self.left._display_aux()
            right, m, q, y = self.right._display_aux()
            s = '%s' % self.val
            u = len(s)
            first_line = (x + 1) * ' ' + (n - x - 1) * \
                '_' + s + y * '_' + (m - y) * ' '
            second_line = x * ' ' + '/' + \
                (n - x - 1 + u + y) * ' ' + '\\' + (m - y - 1) * ' '
            if p < q:
                left += [n * ' '] * (q - p)
            elif q < p:
                right += [m * ' '] * (p - q)
            zipped_lines = zip(left, right)
            lines = [first_line, second_line] + \
                [a + u * ' ' + b for a, b in zipped_lines]
            return lines, n + m + u, max(p, q) + 2, n + u // 2
    
    
    def array_to_tree(arr):
        # Handle case of empty array
        if not arr:
            return None
    
        # Initialize root node and a queue
        root = TreeNode(arr[0] if arr[0] != 'x' else None)
        queue = [root]
    
        # Initialize an index to traverse the array
        index = 1
    
        while queue and index < len(arr):
            node = queue.pop(0)
    
            # Handle left child
            if arr[index] != 'x':
                """
                If the value is 'x', then the left child is None.
                We know this is the left node because we are following a level-order traversal.
                """
                node.left = TreeNode(arr[index])
                """
                Add the left child to the queue so we can handle its children later.
                """
                queue.append(node.left)
            """
            We must increment the index regardless of whether we handled the left child otherwise we will get stuck in an infinite loop.
            """
            index += 1
    
            # Handle right child if there are still nodes left
            if index < len(arr) and arr[index] != 'x':
                """
                If the index is less than the length of the array and the value is 'x', then the right child is None according to level-order traversal.
                """
                node.right = TreeNode(arr[index])
                """
                Just as we did with the left child, we add the right child to the queue so we can handle its children later.
                """
                queue.append(node.right)
            """
            We must increment the index regardless of whether we handled the right child otherwise we will get stuck in an infinite loop.
            """
            index += 1
    
        """
        Note: We intentionally increment the index twice since we are handling two children at a time.
        """
    
        return root
    

    So this:

    if __name__ == "__main__":
        arr = ["6", "x", "8", "4", "x"]
        root = array_to_tree(arr)
    
        # xref: https://stackoverflow.com/questions/34012886/print-binary-tree-level-by-level-in-python
        root.display()
    

    Would output this:

    6_ 
      \
      8
     / 
     4