从二叉树遍历到回溯算法

2356 字

本篇的主题是回溯算法，共有两道例题：

给定一个二叉树和一个目标和，找到所有从根结点到叶结点的路径，使得路径上所有结点值相加等于目标和。

示例：

给定如下二叉树，以及目标和 sum = 22。

返回：
[
 [5,4,11,2],
 [5,8,4,5]
]

给定一组不含重复元素的整数数组 nums，返回该数组所有可能的子集（幂集）。

示例：

输入： nums = [1,2,3]

输出：
[
    []
    [1],
    [2],
    [3],
    [1,2],
    [1,3],
    [2,3],
    [1,2,3],
]

你可能会有疑问，上一篇文章中讲过的 Path Sum 问题为何再次出现？讲解回溯算法为何使用二叉树题目？

实际上，回溯法的求解过程和树的遍历息息相关。在本篇中，我们将从遍历的视角看 Path Sum 问题，由此引出回溯算法的基本思想，并以此思想解决一个正宗的回溯法题目 Subsets。

这篇文章将会包含：

二叉树的遍历过程：以 Path Sum 为例
回溯算法的递归与遍历策略
回溯算法例题：Subsets
相关题目

从遍历的视角看 Path Sum 问题

我们刚刚在上一篇文章中讲过例题 Path Sum，而本篇的例题 Path Sum II 是它的一个变种。两者题目几乎一样，只是要求的输出略有不同。原本只需输出路径的个数，现在要输出所有可能的路径。

而这样一个题目要求的小小变化也让我们看问题的视角发生了变化。在 Path Sum 问题中，我们是从子问题的视角来看二叉树问题的。

从子问题的视角看 Path Sum 问题

而在 Path Sum II 中，为了得到路径，我们需要从遍历的视角来看二叉树问题。

从遍历的角度看 Path Sum 问题

为了最终输出所有可能的路径，我们需要在遍历时记录当前路径，当发现路径满足条件时，就将路径保存下来。

路径遍历的顺序实际上就是 DFS 的顺序。当 DFS 进入一个结点时，路径中就增加一个结点；当 DFS 从一个结点退出时，路径中就减少一个结点。下面的 GIF 动图展示了这个过程。

DFS 中的路径变化（动图）

而 Path Sum 问题，就是在遍历到叶结点，路径最长的时候，把需要的路径保存下来。

回溯算法的基本思想

这样一个二叉树遍历问题如何跟回溯法联系起来呢？我们看看回溯算法的定义：

回溯法采用试错的思想，当它通过尝试发现现有的分步答案不能得到有效的正确的解答的时候，它将取消上一步甚至是上几步的计算，再通过其它的可能的分步解答再次尝试寻找问题的答案。—— 回溯法 - 维基百科

从字面意思上来看，回溯（backtracking） 实际上就是“撤回一步”的意思。而在二叉树的 DFS 遍历中，从一个结点退出就是一种回溯。回溯法和 DFS 是息息相关的。例如，使用 DFS 遍历这个二叉树中的路径，在走到结点 7 之后，下一步就是退回一个结点（递归函数返回），再进入结点 2。这样的一个退回步骤就是回溯。

DFS 中递归函数返回就是一种回溯

根据回溯操作的特性，我们使用栈记录遍历时的当前路径。当进入一个结点时，做 push 操作；当退出一个结点时，做 pop 操作，进行回溯。

语言小贴士：在不同的语言中该用什么数据结构表示栈？

在 C++ 中，一般用 vector 。
vector<int> path;
path.push_back(13);
path.pop_back();
在 Java 中，很多人用 ArrayList。但 ArrayList 没有方便的 pop 操作，所以我更推荐 ArrayDeque。
Deque<Integer> path = new ArrayDeque<>();
path.addLast(13);
path.removeLast();
在 Python 中，直接使用列表即可。
path = []
path.append(13)
path.pop()

最终我们得到的题解代码是：

public List<List<Integer>> pathSum(TreeNode root, int sum) {
    List<List<Integer>> res = new ArrayList<>();
    Deque<Integer> path = new ArrayDeque<>();
    traverse(root, sum, path, res);
    return res;
}

void traverse(TreeNode root, int sum, Deque<Integer> path, List<List<Integer>> res) {
    if (root == null) {
        return;
    }
    path.addLast(root.val);
    if (root.left == null && root.right == null) {
        if (root.val == sum) {
            res.add(new ArrayList<>(path));
        }
    }
    int target = sum - root.val;
    traverse(root.left, target, path, res);
    traverse(root.right, target, path, res);
    path.removeLast();
}

代码的整体结构和上期例题题解类似，只是加上了栈 path 记录当前路径。关于栈的 push 和 pop 操作，有两个需要注意的地方：

保证刚进入结点就 push，最后退出结点之前才 pop，这样才能使当前路径和遍历的进度对应；
在叶结点判断后，不能进行 return，否则会跳过后面的 pop 操作而出错。

这两点都需要做题来体验，建议亲自做一遍例题来体会。

Subsets 问题的遍历方式

上面讲解了一道二叉树题目中的回溯思想，可能你还是觉得不太像回溯法。那么我们再来看一道典型的回溯法问题：Subsets。

Subsets 问题就是要枚举出集合的所有子集。生成子集有一个很简单的策略，一个子集可以选择使用或不使用第一个元素，选好之后，再对第二个元素进行选择，以此类推。这就是一种回溯的思想。

每个元素有两种选择， $n$ 个元素共有 $2^n$ 种可能的选择，对应 $2^n$ 个子集。我们可以将选择的过程画下来。下图是集合 $\{1, 2, 3\}$ 的选择过程，共有 8 个结果。

Subsets 问题的搜索树

这又是一个树的结构。一般来说，回溯算法都可以将决策路径画成树的形状，成为一棵搜索树。回溯法执行的过程实际上就是在这棵树上做遍历。刚好这还是一棵二叉树，这又联系上了二叉树的遍历。

那么，我们可以尝试用遍历树的思路写出回溯法的代码。这里的栈是当前子集里的元素，push 操作是往子集里加元素（上图中的“选1”、“选2”、“选3”），pop 操作是从子集中删除元素（撤销选择）。

最终我们得到完整的代码：

public List<List<Integer>> subsets(int[] nums) {
    Deque<Integer> current = new ArrayDeque<>(nums.length);
    List<List<Integer>> res = new ArrayList<>();
    backtrack(nums, 0, current, res);
    return res;
}

void backtrack(int[] nums, int k, Deque<Integer> current, List<List<Integer>> res) {
    if (k == nums.length) {
        res.add(new ArrayList<>(current));
        return;
    }

    // 不选择第 k 个元素
    backtrack(nums, k+1, current, res);

    // 选择第 k 个元素
    current.addLast(nums[k]);
    backtrack(nums, k+1, current, res);
    // 撤销选择
    current.removeLast();
}

这份代码看起来和 Path Sum II 的代码非常类似，例如都使用了一个栈，递归的参数也很像。但是递归调用和 push/pop 的操作方式有一些微妙的地方。

现在，我们是在调用递归函数之前和之后进行 push/pop，这是因为数组本身并没有递归结构，我们需要用 push/pop 操作来营造出不同的选择。两个递归函数的调用其实都是一样的，但因为 current 中的内容不一样，所以其实是两个决策路径。

回溯算法的复杂度一般都会很高。以 Subsets 问题为例，从搜索树的规模可以看出算法的时间复杂度是非常高的 $O(2^n)$ 。不过，回溯法写成这样的复杂度是可接受的，一般的回溯法题目也没有更高效的解法。实际上，这份代码在用时上战胜了 100% 的提交。

提交结果

总结

通过这两个例题我们看到了回溯算法和二叉树遍历的相似关系。在求解回溯算法的时候，我们可以先构造一个搜索树，在这个树上遍历进行递归求解。

需要注意的是，例题 Subsets 中的搜索树是二叉树，这只是个巧合。实际上搜索树完全可以是多叉树，而且多叉树才更常见。

本篇讲解的是比较基础的回溯法思想。回溯法还有很多技巧，例如 Permutation 和 Combination 系列题目，后续还会有文章进行讲解。

从二叉树遍历到回溯算法

从遍历的视角看 Path Sum 问题

回溯算法的基本思想

Subsets 问题的遍历方式

总结

相关题目