Depth-first search to find the maximum depth of a binary tree

Question

I'm solving a leetcode problem to find a maximum depth of a binary tree. The recursive solution is fairly straightforward so I tried to implement iterative DFS approach. Here is what I come up with:

public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    public TreeNode() {}
    public TreeNode(int val) { this.val = val; }
    public TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

public static int maxDepthDfs(TreeNode root){
    Deque<TreeNode> stack = new LinkedList<>();
    Deque<Integer> depths = new LinkedList<>();
    TreeNode current = root;
    int maxDepth = 0;
    int currentDepth = 0;
    while(!stack.isEmpty() || current != null){
        if(current == null){
            if(currentDepth > maxDepth){
                maxDepth = currentDepth;
            }
            current = stack.pop().right;
            currentDepth = depths.pop() + 1;
        } else {
            stack.push(current);
            current = current.left;
            depths.push(currentDepth);
            currentDepth++;
        }
    }
    return maxDepth;
}

The problem about this solution is that there are to much of extra space. Is there a way to optimize it? Maybe the idea of Morris traversal may be helpful here?

Answer 1

There isn't much space being "wasted" here in my opinion, not sure why you are trying to optimise this. An approach you might try is, add a depth field to the TreeNode class. That should remove the need for using the depths arrays, and save you from calling push/pop operations on the extra array.