lc 1029. recover a tree from preorder traversal

We run a preorder depth first search on the root of a binary tree.
At each node in this traversal, we output D dashes (where D is the depth of this node), then we output the value of this node. (If the depth of a node is D, the depth of its immediate child is D+1. The depth of the root node is 0.)
If a node has only one child, that child is guaranteed to be the left child.
Given the output S of this traversal, recover the tree and return its root.

Example 1:
Input: “1-2–3–4-5–6–7”
Output: [1,2,5,3,4,6,7]

Example 2:
Input: “1-2–3—4-5–6—7”
Output: [1,2,5,3,null,6,null,4,null,7]

We can use a stack to deal with this problem. Traverse the string S, and once we meet a number, advance i until we meet a “-“ and vice versa to get a level. Once the length of stack is larger than current level, that is to say we should pop the stack to get the correct parent of current node, and then set the node as the child of stack[-1].

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

class :
    def recoverFromPreorder(self, S):
        stack, i = [], 0
        while i < len(S):
            level = 0
            val = ""
            while i < len(S) and S[i] == '-':
                level += 1
                i += 1
            while i < len(S) and S[i] != '-':
                val += S[i]
                i += 1
            while len(stack) > level:
                stack.pop()
            node = TreeNode(val)
            if stack and not stack[-1].left:
                stack[-1].left = node
            elif stack:
                stack[-1].right = node
            stack.append(node)
        return stack[0]