Problem overview
A binary tree is written in preorder: a node, then its whole left subtree, then its whole right subtree. Each node is written as D dashes, where D is its depth, followed by its value. The root has depth 0 and no dashes. If a node has only one child, that child is the left child.
The function recoverFromPreorder receives such a string, traversal, and returns the root of the rebuilt tree. It has 1 to 1,000 nodes, and every value is between 1 and 1,000,000,000, so a value can have several digits.
For "1-2--3--4-5--6--7" the root is 1. The nodes with one dash, 2 and 5, are its children, 2 being the left one. The nodes --3 and --4 follow 2, so they are its children, and --6 and --7 follow 5. The level-order output is [1,2,5,3,4,6,7].
Approach: keep the path from the root in a stack
Preorder always continues from the deepest open node. When the next node has depth d, its parent is the last node you saw at depth d - 1. So keep a stack that holds the path from the root down to the most recent node: the node at position d of the stack has depth d.
For each new node, read its dashes to get d, then read its digits to get the value. Next, pop the stack until it holds exactly d nodes. The nodes you drop are finished, since their subtrees are complete, and they stay attached to their own parents. Now the top of the stack is the parent. If the parent's left child is free, the new node goes there; otherwise it goes to the right. Finally, push the new node, since it is now the deepest node of the path.
The depth can jump back by several levels at once. In "1-2--3---4-5--6---7", after ---4 at depth 3 comes -5 at depth 1. Popping until the stack holds one node removes 4, 3 and 2, leaving 1 as the parent of 5. Since 1 already has the left child 2, 5 becomes its right child.
The string is read once, so the time is linear in its length.
Step by step
- Create
stack, an empty slice of*TreeNode. - Create
i, the reading position, at 0. All loops move it forward and never back. - Loop while
i < len(traversal). Each pass reads one node: its dashes, then its value. - Create
depthat 0, to count the dashes. - Loop while
traversal[i]is'-'. No bounds check is needed here, because a value always follows the dashes. - Add one to
depthand advancei. When the loop ends,iis on the first digit. - Create
valueat 0, to build the number digit by digit. - Loop while
iis inside the string and the character is not a dash. The value ends at the next node's dashes or at the end of the string. - Append the digit: multiply
valueby 10, addint(traversal[i]-'0'), and advancei. - Create the node with
&TreeNode{Val: value}. - Drop deeper nodes: loop while the stack holds more than
depthnodes. - Pop the top node in that loop by re-slicing.
- Check that the stack is not empty. Only the root finds it empty.
- Name the top
parent. It stays on the stack, since it may get a second child. - If
parent.Leftisnil, set it to the new node. In preorder the left subtree comes first, so a first child is a left child, and a single child is on the left as the statement says. - Otherwise set
parent.Right. A parent has at most two children. - Push the node with
append, storing the result back instack. - Return
stack[0]. The root is the first node read and is never popped, because every later node has depth 1 or more.
The solution in Go
func recoverFromPreorder(traversal string) *TreeNode {
stack := []*TreeNode{}
i := 0
for i < len(traversal) {
depth := 0
for traversal[i] == '-' {
depth++
i++
}
value := 0
for i < len(traversal) && traversal[i] != '-' {
value = value*10 + int(traversal[i]-'0')
i++
}
node := &TreeNode{Val: value}
for len(stack) > depth {
stack = stack[:len(stack)-1]
}
if len(stack) > 0 {
parent := stack[len(stack)-1]
if parent.Left == nil {
parent.Left = node
} else {
parent.Right = node
}
}
stack = append(stack, node)
}
return stack[0]
}
Go details worth a closer look:
- Bytes from a string.
traversal[i]is abyte. It compares directly with the rune constant'-', because an untyped character constant adapts to the byte type. - Digit to number.
traversal[i]-'0'subtracts the code of'0'(48) from the digit's code, giving 0 to 9 as abyte.int(...)converts it so that it can be added to theintvariablevalue. Go does not mix types in arithmetic. - Pointers to nodes.
&TreeNode{Val: value}allocates a node and gives its address. Settingparent.Left = nodestores that pointer, and later changes to the node are visible through both names. - The loop that pops.
stack = stack[:len(stack)-1]shortens the slice with no copying. The popped nodes remain in the tree through their parents. - Short-circuit order. In
i < len(traversal) && traversal[i] != '-'the length test comes first, so the index is never out of range. The dash loop has no such test, which is safe only because the input is valid and never ends in a dash.
Complexity
Time: O(n), where n is the length of the string. Each character is read once. Each node is pushed once and popped at most once.
Space: O(h) for the stack, where h is the height of the tree, plus the O(m) nodes of the tree being built, with m nodes.
Edge cases and common mistakes
- Multi-digit values. In
"5-12--7---40--9"the digits1and2form 12, and4and0form 40. Reading one character per node would break these. - A single node.
"7"gives a tree with only the root. The stack is empty when the root is read, so it gets no parent. - Single children go left. In
"1-2-3"nodes 2 and 3 are both children of 1, so 2 is on the left and 3 on the right. In"1-2--3", node 3 is the only child of 2 and goes on the left. - Jumping several levels up. The pop loop removes as many nodes as needed. Popping only one would attach the node to the wrong parent.
- Popping the parent. The loop keeps exactly
depthnodes, so the parent, at positiondepth - 1, stays.
Testing
"1-2--3--4-5--6--7"gives[1,2,5,3,4,6,7]."1-2--3---4-5--6---7"gives[1,2,5,3,null,6,null,4,null,7]."5-12--7---40--9"gives[5,12,null,7,9,40]."7"gives[7]."1-2-3"gives[1,2,3].
