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;
}