(FB) 从 Binary Tree Path 看如何递归转迭代

例题：给定 Binary Tree，print 所有 root - leaf 的 path.

          A
         / \
        B   D
       /   / \
      C   E   F

• Input : Root A

• Output :

  • A,B,C

  • A,D,E

  • A,D,F

递归：

比较简单直接，就是在树上 DFS + backtrack，带着当前节点 root 和当前路径 list 递归处理即可。这题的结构和 LC 那道 Binary Paths 一样。

    public static void printPaths_recursion(TreeNode root){
        dfs(root, new ArrayList<>());
    }

    private static void dfs(TreeNode root, List<TreeNode> list){
        if(root == null) return;

        list.add(root);

        if(root.left == null && root.right == null){
            for(TreeNode node : list) System.out.print(node.val + ",");
            System.out.println();
            list.remove(list.size() - 1);
            return;
        }

        dfs(root.left, list);
        dfs(root.right, list);

        list.remove(list.size() - 1);
    }

迭代解法 1：

• Queue，空间占用和时间都高一些。

• 基本思路就是用一个 Queue 存所有当前的 path，队列中的 path 数量与进出关系都和正常的 level order traversal 完全一样，每次队列末尾的节点，都是 current node.

• 中间 poll 的时候看到叶子节点，就把 list 输出出来就好。

• 这种解法的问题一是空间占用是 O(N) 的，相比 Stack 高；二是每次生成新 node 的时候其实是在做一次 collection 的 copy 操作，时间复杂度上也不经济。

  private static void printPath_bfs(TreeNode root){
      Queue<List<TreeNode>> queue = new LinkedList<>();
  
      if(root != null){
          List<TreeNode> rootList = new ArrayList<>();
          rootList.add(root);
          queue.offer(rootList);
      }
  
      while(!queue.isEmpty()){
          List<TreeNode> curList = queue.poll();
          TreeNode curNode = curList.get(curList.size() - 1);
  
          if(curNode.left == null && curNode.right == null){
              for(TreeNode node : curList) System.out.print(node.val + ",");
              System.out.println();
          }
  
          if(curNode.left != null){
              List<TreeNode> list = new ArrayList<>(curList);
              list.add(curNode.left);
              queue.offer(list);
          }
          if(curNode.right != null){
              List<TreeNode> list = new ArrayList<>(curList);
              list.add(curNode.right);
              queue.offer(list);
          }
      }
  }

迭代解法 2:

Stack + HashSet

迭代解法 3: Stack + 自定义 StackFrame

这种做法最有意思，我也最喜欢，觉得只要有合适的 field 属性，很容易 generalize 到各种其他递归解法中。

首先对于二叉树，我们可以这样定义 StackFrame 的三个状态：

• 0 : 刚入栈，未访问子树；

• 1 : 正在访问左子树，返回代表左子树访问完毕；

• 2 : 正在访问右子树，返回代表右子树访问完毕；

这样代码如下，每次 stack 里存的，都是当前路径 (配合一个 List 做高效 print path 操作)，而每次栈顶的 StackFrame 都记录了当前 node 的访问状态和下一步的动向。

这种 0/1/2 的状态表示方式看着有点像 Graph 里面做 DFS / BFS 的标注，其实不完全一样，只是在这题我们的树一定是二叉而已。

对于多叉树的情况，改动也不会很难；这次 index 代表着 “下一个要访问的 child index”，当 index == children.size() 的时候，我们就可以知道当前 node 的所有子节点都访问完了。

public static void printPaths(TreeNode root){
        Stack<StackFrame> stack = new Stack<>();
        List<TreeNode> list = new ArrayList<>();
        if(root != null) stack.push(new StackFrame(0, root));

        while(!stack.isEmpty()){
            StackFrame curFrame = stack.peek();
            TreeNode curNode = curFrame.node;

            if(curNode == null) {
                stack.pop();
            } else if(curFrame.status == 0){
                list.add(curNode);
                curFrame.status = 1;
                stack.push(new StackFrame(0, curNode.left));
            } else if(curFrame.status == 1){
                curFrame.status = 2;
                stack.push(new StackFrame(0, curNode.right));
            } else {
                if(curNode.left == null && curNode.right == null) {
                    for (TreeNode node : list) System.out.print("" + node.val + ",");
                    System.out.println();
                }
                stack.pop();
                list.remove(list.size() - 1);
            }
        }
    }