Problem overview
The method zigzagTraversal receives the root of a binary tree of TreeNode objects and returns the node values as an Array[Array[Int]], one inner array per level. The first level (the root) reads left to right, the second right to left, the third left to right again, and so on.
For the tree with root 3, children 9 and 20, and 20's children 15 and 7, the answer is [[3], [20, 9], [15, 7]]. An empty tree (a null root) gives an empty array.
Approach: a stack reverses the order for free
A plain level order traversal reads each level left to right. One way to get the zigzag is to collect each level normally and reverse every second one. This lesson avoids the reversing step by using the way a stack works: the last element pushed is the first one popped.
The traversal keeps two stacks. currentLevelStack holds the nodes of the level being read, and nextLevelStack collects their children. Nodes are popped from the current stack, their values recorded, and their children pushed onto the other stack. When the current stack is empty, the level is finished, and the other stack becomes the current one.
The direction of the next level depends on the order in which the children are pushed:
- On an even level (read left to right), push the left child first and the right child second. The right child ends on top and is popped first, so the next level reads right to left.
- On an odd level (read right to left), push the right child first and the left child second. The left child is popped first, so the next level reads left to right.
Here is the example. Level 0 holds [3]. Popping 3 records it and pushes 9, then 20, so the next stack is [9, 20] and 20 is on top. Level 1 pops 20 and then 9, which gives [20, 9]. Level 1 is odd, so 20 pushes its children as 7, then 15, and 9 pushes nothing. The next stack is [7, 15], and level 2 pops 15 and then 7, which gives [15, 7].
Step by step
- Create
resultas a mutable buffer of mutable buffers ofInt. Each inner buffer will hold one level in its final order. The code spells out the full package namescala.collection.mutable.ArrayBufferinstead of importing it. - Create the two stacks as empty
ArrayBuffer[TreeNode]values. They are declared withvarbecause the traversal reassigns them when a level ends. AnArrayBufferworks as a stack when you only add and remove at its end. - Push the root through the helper
addToStack. The helper skipsnull, so an empty tree leaves the stack empty and the loop never runs. - Track the level number in
level, starting at 0. Its parity decides the push order. - Loop while
currentLevelStack.nonEmpty. Each iteration handles one node, and the loop ends when the last level has been read. - Pop a node with
remove(currentLevelStack.length - 1), which removes the last element and returns it. - Record its value with
addToResult, which creates the inner buffer for a level the first time a node of that level shows up. - Push the children in the order that depends on
level % 2: left then right on even levels, right then left on odd ones. Missing children are passed in too, and the helper filters them out. - Detect the end of a level. After the push, an empty
currentLevelStackmeans the node just handled was the last of its level. The next stack then holds the whole next level, so it becomes the current stack, a fresh empty buffer takes its place, andlevelgoes up by one. - Convert the result to
Array[Array[Int]]withresult.map(_.toArray).toArray. - Write the helpers
addToResultandaddToStackabovezigzagTraversal.
The solution in Scala
class Solution {
def addToStack(stack: scala.collection.mutable.ArrayBuffer[TreeNode], nodes: Array[TreeNode]): Unit = {
for (node <- nodes) {
if (node != null) {
stack += node
}
}
}
def addToResult(result: scala.collection.mutable.ArrayBuffer[scala.collection.mutable.ArrayBuffer[Int]], level: Int, value: Int): Unit = {
if (level >= result.length) {
result += scala.collection.mutable.ArrayBuffer[Int]()
}
result(level) += value
}
def zigzagTraversal(root: TreeNode): Array[Array[Int]] = {
val result = scala.collection.mutable.ArrayBuffer[scala.collection.mutable.ArrayBuffer[Int]]()
var currentLevelStack = scala.collection.mutable.ArrayBuffer[TreeNode]()
var nextLevelStack = scala.collection.mutable.ArrayBuffer[TreeNode]()
addToStack(currentLevelStack, Array(root))
var level = 0
while (currentLevelStack.nonEmpty) {
val node = currentLevelStack.remove(currentLevelStack.length - 1)
addToResult(result, level, node.value)
if (level % 2 == 0) {
addToStack(nextLevelStack, Array(node.left, node.right))
} else {
addToStack(nextLevelStack, Array(node.right, node.left))
}
if (currentLevelStack.isEmpty) {
currentLevelStack = nextLevelStack
nextLevelStack = scala.collection.mutable.ArrayBuffer[TreeNode]()
level += 1
}
}
result.map(_.toArray).toArray
}
}
The Scala details worth a closer look:
- The last expression is the result.
result.map(_.toArray).toArrayhas noreturn. The innermapturns eachArrayBuffer[Int]into anArray[Int], and the outertoArrayturns the buffer of arrays into anArray[Array[Int]]. valfixes the name, not the contents.resultis aval, yetresult += ...andresult(level) += valuestill change the buffers. The two stacks arevarbecause the names themselves are reassigned.- Assigning a buffer shares it.
currentLevelStack = nextLevelStackmakes both names point at the same buffer. That is why the next line givesnextLevelStacka brand-new buffer. Without it, the children of the next level would be pushed onto the stack that is being popped. remove(index)returns the element. On anArrayBuffer, removing the last element is cheap, so this is a constant-time pop. The buffer does not shrink its backing array, but that doesn't matter here.result(level)reads,+=appends. Round brackets index a buffer. Becauseresult(level)is itself a buffer,result(level) += valueappends to that inner buffer.Array(node.left, node.right)may containnull. A missing child is anullreference in this lesson'sTreeNode, which is whyaddToStacktestsnode != nullbefore pushing.
Complexity
Time: O(n) for n nodes. Each node is pushed once, popped once and recorded once, and every one of those operations takes constant time.
Space: O(n). The result holds every value. Besides it, the two stacks together hold at most the nodes of two neighbouring levels, which is at most n in the worst case, for example in a complete tree where the bottom level has about half of all nodes.
Edge cases and common mistakes
- Empty tree.
zigzagTraversal(null)returns an empty array, because the helper never pushes thenullroot. - Single node. The result is
[[v]], with one inner array. - Chain-shaped trees. A tree where every node has only a left child, such as 1, 2, 3, 4 going down, gives
[[1], [2], [3], [4]]. Each level has one node, so the direction makes no difference. - Full tree. For root 1 with children 2 and 3, and leaves 4, 5, 6, 7 below, the output is
[[1], [3, 2], [4, 5, 6, 7]]. - Pushing in the wrong order. Swapping the two push orders flips every level, so level 0 is still correct but level 1 reads left to right.
- Forgetting the new buffer. If
nextLevelStackis not replaced after the swap, both names refer to the same buffer, and the traversal mixes two levels. - Changing
leveltoo early. The level number must change only when the current stack is empty. Incrementing it for each node would put every value in a different inner array.
Testing
- Root 3 with children 9 and 20, and 20's children 15 and 7, gives
[[3], [20, 9], [15, 7]]. - A
nullroot gives[]. - Root 1 with children 2 and 3, and leaves 4, 5, 6, 7 gives
[[1], [3, 2], [4, 5, 6, 7]]. - A downward chain 1, 2, 3, 4 on the left gives
[[1], [2], [3], [4]].
