如何根据这样的思路写出代码呢?当然可以用 HashSet 这样的数据结构来表示候选集合。但如果你这么去写的话,会发现代码写起来比较啰嗦,而且 set 结构的“遍历-删除”操作不太好写。在这里,我不展示使用 set 结构的代码。大家只要明白一条要点:**在一般情况下,候选集合使用数组表示即可。**候选集合上需要做的操作并不是很多,使用数组简单又高效。
在子集问题中,我们定义了变量 k,表示当前要对第 k 个元素做决策。实际上,变量 k 就是候选集合的边界,指针 k 之后的元素都是候选元素,而 k 之前都是无效元素,不可以再选了。
用数组表示候选集合
而每次决策完之后将 k 加一,就是将第 k 个元素移出了候选集合。
将第 k 个元素移出候选集合
在全排列问题中,我们要处理的情况更难一些。每次做决策时,候选集合中的所有元素都可以选择,也就是有可能删除候选集合中间的元素,这样数组中会出现“空洞”。这种情况该怎么处理呢?我们可以使用一个巧妙的方法,先将要删除的元素与第 k 个元素交换,再将 k 加一,过程如下方动图所示:
理解了图中的关系之后,题解代码就呼之欲出了。我们只需使用一个 current 数组,左半边表示已选元素,右半边表示候选元素。指针 k 不仅是候选元素的开始位置,还是已选元素的结束位置。我们可以得一份非常简洁的题解代码:
java
public List<List<Integer>> permute(List<Integer> nums) { List<Integer> current = new ArrayList<>(nums); List<List<Integer>> res = new ArrayList<>(); backtrack(current, 0, res); return res;}// current[0..k) 是已选集合, current[k..N) 是候选集合void backtrack(List<Integer> current, int k, List<List<Integer>> res) { if (k == current.size()) { res.add(new ArrayList<>(current)); return; } // 从候选集合中选择 for (int i = k; i < current.size(); i++) { // 选择数字 current[i] Collections.swap(current, k, i); // 将 k 加一 backtrack(current, k+1, res); // 撤销选择 Collections.swap(current, k, i); }}
注意写在递归函数上方的注释。在写回溯法问题的代码时,你需要时刻清楚什么是已选集合,什么是候选集合。注释中的条件叫做“不变式”。一方面,我们在函数中可以参考变量 k 的含义,另一方面,我们在做递归调用的时候,要保证这个条件始终成立。特别注意代码中递归调用传入的参数是 k+1 ,即删除一个候选元素,而不是传入 i+1。
n 中取 k 的排列
全排列问题是 n 中取 n 的排列,可以记为 P(n,n)。在面试中,我们很可能会遇到各种各样的变种题,那么 n 中取 k 的排列 P(n,k)、组合 C(n,k) 我们也要掌握。
P(n,k) 问题非常简单,我们只需要在全排列的基础上,做完第 k 个决策后就将结果返回。也就是说,只遍历决策树的前 k 层。例如 n=3,k=2 时,决策树的第 2 层,已选集合中有两个元素,将这里的结果返回即可。
n 中取 k 的排列的决策树
题解代码如下所示,只需要修改递归结束的条件即可。
java
public List<List<Integer>> permute(List<Integer> nums, int k) { List<Integer> current = new ArrayList<>(nums); List<List<Integer>> res = new ArrayList<>(); backtrack(k, current, 0, res); return res;}// current[0..m) 是已选集合, current[m..N) 是候选集合void backtrack(int k, List<Integer> current, int m, List<List<Integer>> res) { // 当已选集合达到 k 个元素时,收集结果并停止选择 if (m == k) { res.add(new ArrayList<>(current.subList(0, k))); return; } // 从候选集合中选择 for (int i = m; i < current.size(); i++) { // 选择数字 current[i] Collections.swap(current, m, i); backtrack(k, current, m+1, res); // 撤销选择 Collections.swap(current, m, i); }}
注意这里 k 是题目的输入,所以原先我们代码里的变量 k 重命名成了 m。此外,就是递归函数开头的 if 语句条件发生了变化,当已选集合达到 k 个元素时,就收集结果停止递归。
组合问题:失效元素
由于排列组合的密切联系,组合问题 C(n,k) ,即 n 中取 k 的组合,可以在 P(n,k) 问题的解法上稍加修改而来。