Problem Statement in English

You’re given a string traversal that represents the preorder traversal of a binary tree. The string consists of digits and dashes.

Each node’s value is represented by a sequence of digits, and the depth of each node is represented by the number of dashes preceding its value. The root node has depth 0, its children have depth 1, and so on.

Return the root of the binary tree represented by traversal.

Approach

There’s 2 parts to this:

  • parse the traversal string to get the depth and value of each node
  • reconstruct the tree

You can do it in 2 separate passes, but you can also do it in a single elegant pass.

For parsing the traversal we’ll just iterate over the string and count dashes. When we realise that the next character is a digit, we can parse the number and create a node.

While the depth of the stack is greater than the depth of the new node, we keep popping.

The top of the stack will always be the most recent node we’ve seen (unless it’s the first node, in which case we just append it to the stack).

If the depth of the new node is greater than the depth of the top of the stack, then it’s a child of that node. If it’s less than or equal, we pop from the stack until we find a node with a depth less than the new node’s depth.

Since it’s preorder, the first child we encounter will always be the left child. If the left child is already filled, then it must be the right child.

When we’re done there will be only one node left in the stack, which is the root of the tree. We can return that node.

And we’re done!

Solution in Python


class Solution:
    def recoverFromPreorder(self, traversal: str) -> Optional[TreeNode]:
        stack = []
        i = 0
        n = len(traversal)

        while i < n:
            # 1. Count dashes to determine depth
            depth = 0
            while i < n and traversal[i] == '-':
                depth += 1
                i += 1

            # 2. Extract digits to get node value
            val = 0
            while i < n and traversal[i].isdigit():
                val = val * 10 + int(traversal[i])
                i += 1

            node = TreeNode(val)

            # 3. Maintain stack depth invariant
            # The top of stack should be the immediate parent (size == depth)
            while len(stack) > depth:
                stack.pop()

            # 4. Attach node to parent
            if stack:
                if not stack[-1].left:
                    stack[-1].left = node
                else:
                    stack[-1].right = node

            stack.append(node)

        return stack[0]

Complexity

  • Time: $O(n)$
    Since we iterate over the traversal string once, where $n$ is the length of the traversal string.

  • Space: $O(n)$
    Since we use a stack to keep track of the nodes, in the worst case, the stack can grow to the size of the number of nodes in the tree.

Mistakes I Made

The version I came up with wasn’t as elegant. The stack wasn’t managed as well as this version.

And we are done.