LeetCode: Binary Tree Zigzag Level Order Traversal C#

Question

https://leetcode.com/explore/interview/card/top-interview-questions-medium/108/trees-and-graphs/787/

Given a binary tree, return the zigzag level order traversal of its nodes' values. (ie, from left to right, then right to left for the next level and alternate between).

For example:

Given binary tree [3,9,20,null,null,15,7],
 3
 / \
 9 20
 / \
 15 7

return its zigzag level order traversal as:

[
 [3],
 [20,9],
 [15,7]
]

/// <summary>
/// https://leetcode.com/explore/interview/card/top-interview-questions-medium/108/trees-and-graphs/787/
/// </summary>
[TestClass]
public class ZigzagLevelOrderTest
{
 [TestMethod]
 public void TestZigZag()
 {
 /* 3
 / \
 9 20
 / \
 15 7
 */
 TreeNode root = new TreeNode(3);
 root.left = new TreeNode(9);
 root.right = new TreeNode(20);
 root.right.left = new TreeNode(15);
 root.right.right = new TreeNode(7);
 IList<IList<int>> res = ZigzagLevelOrderClass.ZigzagLevelOrder(root);
 CollectionAssert.AreEqual(new List<int> { 3 }, res[0].ToList());
 CollectionAssert.AreEqual(new List<int> { 20, 9 }, res[1].ToList());
 CollectionAssert.AreEqual(new List<int> { 15, 7 }, res[2].ToList());
 }
 [TestMethod]
 public void FailedTest()
 {
 /* 1
 / \
 2 3
 / \
 4 5
 */
 TreeNode root = new TreeNode(1);
 root.left = new TreeNode(2);
 root.right = new TreeNode(3);
 root.left.left = new TreeNode(4);
 root.right.right = new TreeNode(5);
 IList<IList<int>> res = ZigzagLevelOrderClass.ZigzagLevelOrder(root);
 CollectionAssert.AreEqual(new List<int> { 1 }, res[0].ToList());
 CollectionAssert.AreEqual(new List<int> { 3, 2 }, res[1].ToList());
 CollectionAssert.AreEqual(new List<int> { 4, 5 }, res[2].ToList());
 }
}
public class ZigzagLevelOrderClass
{
 public static IList<IList<int>> ZigzagLevelOrder(TreeNode root)
 {
 List<IList<int>> result = new List<IList<int>>();
 if (root == null)
 {
 return result;
 }
 Stack<TreeNode> currentLeveL = new Stack<TreeNode>();
 Stack<TreeNode> nextLevel = new Stack<TreeNode>();
 currentLeveL.Push(root);
 while (currentLeveL.Count > 0 || nextLevel.Count > 0)
 {
 var nodes = new List<int>();
 while (currentLeveL.Count > 0)
 {
 var curr = currentLeveL.Pop();
 if (curr.left != null)
 {
 nextLevel.Push(curr.left);
 }
 if (curr.right != null)
 {
 nextLevel.Push(curr.right);
 }
 nodes.Add(curr.val);
 }
 if (nodes.Count > 0)
 {
 result.Add(nodes);
 }
 nodes = new List<int>();
 while (nextLevel.Count > 0)
 {
 var curr = nextLevel.Pop();
 if (curr.right != null)
 {
 currentLeveL.Push(curr.right);
 }
 if (curr.left != null)
 {
 currentLeveL.Push(curr.left);
 }
 nodes.Add(curr.val);
 }
 if (nodes.Count > 0)
 {
 result.Add(nodes);
 }
 }
 return result;
 }
}

Please review for performance. this is the second time I solved this question. I think I did a better job now.

Answer 1

Review

• I would return IEnumerable<IEnumerable<int>> rather than IList<IList<int>>. We don't want the caller to change the return value, only to iterate it.
• The two inner loops are almost exactly the same, except that the order of node.left and node.right gets swapped. This part I would refactor to get DRY code.
• You should use var a bit more often: Stack<TreeNode> currentLeveL = new Stack<TreeNode>(); -> var currentLevel = new Stack<TreeNode>(); (also notice the small casing typo in currentLeveL)
• if (root == null) return result; -> perhaps the challenge specifies this edge case, but I would prefer an ArgumentNullException when the input is null and clearly shouldn't be.

---

Refactored

• We can avoid using an outer loop with two nearly identical inner loops, if we exchange currentLevel for nextLevel after each inner loop and use a bool zig that toggles for every cycle of the inner loop to get the zig-zag effect.
• Notice I made an instance method rather than extension method, but feel free to keep an extension method instead. I find traversal to be part of the instance operations.
• I expect performance to remain the same. We're still using two stacks the same way.

public IEnumerable<IEnumerable<int>> ZigzagLevelOrder()
{
 var levels = new List<IEnumerable<int>>();
 var currentLevel = new Stack<TreeNode>();
 var nextLevel = new Stack<TreeNode>();
 var zig = false;
 currentLevel.Push(this);
 while (currentLevel.Any())
 {
 levels.Add(currentLevel.Select(n => n.Value).ToArray());
 zig = !zig;
 while (currentLevel.Any())
 {
 var node = currentLevel.Pop();
 if (zig && node.left != null)
 nextLevel.Push(node.left);
 if (node.right != null)
 nextLevel.Push(node.right);
 if (!zig && node.left != null)
 nextLevel.Push(node.left);
 }
 currentLevel = nextLevel;
 nextLevel = new Stack<TreeNode>();
 }
 return levels.ToArray();
}

And the unit tests pass:

[TestMethod]
public void Fixture()
{
 var root = new TreeNode(3);
 root.left = new TreeNode(9);
 root.right = new TreeNode(20);
 root.right.left = new TreeNode(15);
 root.right.right = new TreeNode(7);
 var res = root.ZigzagLevelOrder();
 CollectionAssert.AreEqual(new List<int> { 3 }, res.First().ToList());
 CollectionAssert.AreEqual(new List<int> { 20, 9 }, res.Skip(1).First().ToList());
 CollectionAssert.AreEqual(new List<int> { 15, 7 }, res.Skip(2).First().ToList());
 root = new TreeNode(1);
 root.left = new TreeNode(2);
 root.right = new TreeNode(3);
 root.left.left = new TreeNode(4);
 root.right.right = new TreeNode(5);
 res = root.ZigzagLevelOrder();
 CollectionAssert.AreEqual(new List<int> { 1 }, res.First().ToList());
 CollectionAssert.AreEqual(new List<int> { 3, 2 }, res.Skip(1).First().ToList());
 CollectionAssert.AreEqual(new List<int> { 4, 5 }, res.Skip(2).First().ToList());
}

---

Performance Optimization

• (削除) I'm not sure how or whether performance could still be optimized beyond the OP code (without falling into a micro-optimisation trap). (削除ここまで)

• In hindsight, after reading through JAD's answer, a further optimisation is to use foreach (var node in currentLevel) rather than while (currentLevel.Any()) to avoid var node = currentLevel.Pop();, which is no longer required (as opposed to OP code) since we exchange the instance of currentLevel with nextLevel anyway.

Answer 2

Your implementation is close, but it can be a bit shorter. You correctly use two stacks, but you duplicate the code alternating between using the two. That's a bit of a waste of space. Instead, at the end of the first while loop, you can just assign nextLevel to currentLevel, create a new stack to nextLevel, and repeat:

 while (currentLeveL.Count > 0 || nextLevel.Count > 0)
 {
 var nodes = new List<int>();
 while (currentLeveL.Count > 0)
 {
 var curr = currentLeveL.Pop();
 if (curr.left != null)
 {
 nextLevel.Push(curr.left);
 }
 if (curr.right != null)
 {
 nextLevel.Push(curr.right);
 }
 nodes.Add(curr.val);
 }
 if (nodes.Count > 0)
 {
 result.Add(nodes);
 }
 currentLevel = nextLevel;
 nextLevel = new Stack<TreeNode>();
 }

The while condition can also be made easier, nextLevel is known to be empty when it is evaluated:

while (currentLeveL.Count > 0)

---

Since you're looping over the stack, only removing each item, not adding items back or anything, you can just replace the while for a foreach:

public static IList<IList<int>> ZigzagLevelOrder(TreeNode root)
{
 List<IList<int>> result = new List<IList<int>>();
 if (root == null)
 {
 return result;
 }
 Stack<TreeNode> currentLeveL = new Stack<TreeNode>();
 Stack<TreeNode> nextLevel = new Stack<TreeNode>();
 currentLeveL.Push(root);
 while (currentLeveL.Count > 0)
 {
 var nodes = new List<int>();
 foreach(var curr in currentLeveL)
 { 
 if (curr.left != null)
 {
 nextLevel.Push(curr.left);
 }
 if (curr.right != null)
 {
 nextLevel.Push(curr.right);
 }
 nodes.Add(curr.val);
 }
 if (nodes.Count > 0)
 {
 result.Add(nodes);
 }
 currentLeveL = nextLevel;
 nextLevel = new Stack<TreeNode>();
 }
 return result;
}

• VisualMelon: Please forget I said anything: I'm clearly not thinking properly this morning! (though you can ditch the nodes.Count > 0 check)
• Gilad: @JAD, your code doesn't pass my unit tests. please notice when you switch between current Level and next level you also need to reverse the order of pushing first right then left node.
• JAD: @gilad hmm, I didn't notice the left/right switching. Let me think.