Leetcode: Binary Tree Paths

Question

I'm solving the following leetcode problem:

Given the root of a binary tree, return all root-to-leaf paths in any order.

A leaf is a node with no children.

Input: root = [1,2,3,null,5]

Output: ["1->2->5","1->3"]

Recursive solution is trivial so I tried to come up with iterative solution. Here is how my solution looks like:

public static List<String> binaryTreePaths(TreeNode root) {
    List<String> retVal = new ArrayList<>();
    Deque<TreeNode> stack = new ArrayDeque<>();
    while(root != null) {
        if(root.left != null){
            stack.push(root);
            root = root.left;
        } else if (root.right != null) {
            stack.push(root);
            root = root.right;
        //
        //The code under the else branch below is difficult to read
        //
        } else {
            final var arr = new ArrayList<>(stack);
            Collections.reverse(arr);
            arr.add(root);
            final var path = arr.stream()
                    .map(node -> String.valueOf(node.val))
                    .collect(Collectors.joining("->"));
            retVal.add(path);
            TreeNode newRoot = null;
            while(!stack.isEmpty()){
                TreeNode previous = stack.pop();
                TreeNode current = root;
                if(previous.right != null && previous.right != current){
                    stack.push(previous);
                    newRoot = previous.right;
                    break;
                } else {
                    root = previous;
                }
            }
            root = newRoot;
        }
    }
    return retVal;
}

The problem: The code is relatively verbose and difficult to read. My idea is to keep track of the whole path from the top to the current node. But this differs from the standard DFS since it keeps only unvisited node on the stack.

Is there a way to improve the part and/or the whole implementation?

Answer 1

I'm having trouble reasoning about the correctness of your code. The moving of root got me lost. Why not build the strings as you go?:

public static List<String> binaryTreePaths(TreeNode root) {
    List<String> retVal = new ArrayList<>();
    Stack<AbstractMap.SimpleEntry<TreeNode, String>> stack = new Stack<>();
    if (root == null) return retVal;
    stack.push(makePair(root, root.val));
    while(!stack.empty()) {
        AbstractMap.SimpleEntry<TreeNode, String> item = stack.pop();
        TreeNode current = item.getKey();
        if(current.left!=null) {
            stack.push(makePair(current.left,
                item.getValue()+"->"+current.left.val));
        }
        if(current.right!=null) {
            stack.push(makePair(current.right,
                item.getValue()+"->"+current.right.val));
        }
        if(current.left==null && current.right==null) {
            retVal.add(item.getValue());
        }
    }
    return retVal;
}

Code is untested and probably won't compile, treat as pseudocode. Implementation of makePair left as exercise to the reader (AbstractMap.SimpleEntry used for a Pair, you can use some other Pair/Tuple).