回溯算法牛逼！
作者：labuladong
公眾號：labuladong
之前說過回溯算法是筆試中最好用的算法，只要你沒什么思路，就用回溯算法暴力求解，即便不能通過所有測試用例，多少能過一點。
回溯算法的技巧也不難，前文回溯算法框架套路 說過，回溯算法就是窮舉一棵決策樹的過程，只要在遞歸之前「做選擇」，在遞歸之后「撤銷選擇」就行了。
但是，就算暴力窮舉，不同的思路也有優(yōu)劣之分。
本文就來看一道非常經典的回溯算法問題，子集劃分問題，可以幫你更深刻理解回溯算法的思維，得心應手地寫出回溯函數。
題目非常簡單：
給你輸入一個數組nums和一個正整數k，請你判斷nums是否能夠被平分為元素和相同的k個子集。
函數簽名如下：
boolean canPartitionKSubsets(int[] nums, int k);
我們之前 背包問題之子集劃分 寫過一次子集劃分問題，不過那道題只需要我們把集合劃分成兩個相等的集合，可以轉化成背包問題用動態(tài)規(guī)劃技巧解決。
但是如果劃分成多個相等的集合，解法一般只能通過暴力窮舉，時間復雜度爆表，是練習回溯算法和遞歸思維的好機會。
一、思路分析
把裝有n個數字的數組nums分成k個和相同的集合，你可以想象將n個數字分配到k個「桶」里，最后這k個「桶」里的數字之和要相同。
前文 回溯算法框架套路 說過，回溯算法的關鍵在哪里？
關鍵是要知道怎么「做選擇」，這樣才能利用遞歸函數進行窮舉。
那么回想我們這個問題，將n個數字分配到k個桶里，我們可以有兩種視角：
視角一，如果我們切換到這n個數字的視角，每個數字都要選擇進入到k個桶中的某一個。
視角二，如果我們切換到這k個桶的視角，對于每個桶，都要遍歷nums中的n個數字，然后選擇是否將當前遍歷到的數字裝進自己這個桶里。
你可能問，這兩種視角有什么不同？
用不同的視角進行窮舉，雖然結果相同，但是解法代碼的邏輯完全不同；對比不同的窮舉視角，可以幫你更深刻地理解回溯算法，我們慢慢道來。
二、以數字的視角
用 for 循環(huán)迭代遍歷nums數組大家肯定都會：
for (int index = 0; index < nums.length; index++) {
System.out.println(nums[index]);
}
遞歸遍歷數組你會不會？其實也很簡單：
void traverse(int[] nums, int index) {
if (index == nums.length) {
return;
}
System.out.println(nums[index]);
traverse(nums, index + 1);
}
只要調用traverse(nums, 0)，和 for 循環(huán)的效果是完全一樣的。
那么回到這道題，以數字的視角，選擇k個桶，用 for 循環(huán)寫出來是下面這樣：
// k 個桶（集合），記錄每個桶裝的數字之和
int[] bucket = new int[k];
// 窮舉 nums 中的每個數字
for (int index = 0; index < nums.length; index++) {
// 窮舉每個桶
for (int i = 0; i < k; i++) {
// nums[index] 選擇是否要進入第 i 個桶
// ...
}
}
如果改成遞歸的形式，就是下面這段代碼邏輯：
// k 個桶（集合），記錄每個桶裝的數字之和
int[] bucket = new int[k];
// 窮舉 nums 中的每個數字
void backtrack(int[] nums, int index) {
// base case
if (index == nums.length) {
return;
}
// 窮舉每個桶
for (int i = 0; i < bucket.length; i++) {
// 選擇裝進第 i 個桶
bucket[i] += nums[index];
// 遞歸窮舉下一個數字的選擇
backtrack(nums, index + 1);
// 撤銷選擇
bucket[i] -= nums[index];
}
}
雖然上述代碼僅僅是窮舉邏輯，還不能解決我們的問題，但是只要略加完善即可：
// 主函數
public boolean canPartitionKSubsets(int[] nums, int k) {
// 排除一些基本情況
if (k > nums.length) return false;
int sum = 0;
for (int v : nums) sum += v;
if (sum % k != 0) return false;
// k 個桶（集合），記錄每個桶裝的數字之和
int[] bucket = new int[k];
// 理論上每個桶（集合）中數字的和
int target = sum / k;
// 窮舉，看看 nums 是否能劃分成 k 個和為 target 的子集
return backtrack(nums, 0, bucket, target);
}
// 遞歸窮舉 nums 中的每個數字
boolean backtrack(
int[] nums, int index, int[] bucket, int target) {
if (index == nums.length) {
// 檢查所有桶的數字之和是否都是 target
for (int i = 0; i < bucket.length; i++) {
if (bucket[i] != target) {
return false;
}
}
// nums 成功平分成 k 個子集
return true;
}
// 窮舉 nums[index] 可能裝入的桶
for (int i = 0; i < bucket.length; i++) {
// 剪枝，桶裝裝滿了
if (bucket[i] + nums[index] > target) {
continue;
}
// 將 nums[index] 裝入 bucket[i]
bucket[i] += nums[index];
// 遞歸窮舉下一個數字的選擇
if (backtrack(nums, index + 1, bucket, target)) {
return true;
}
// 撤銷選擇
bucket[i] -= nums[index];
}
// nums[index] 裝入哪個桶都不行
return false;
}
有之前的鋪墊，相信這段代碼是比較容易理解的。這個解法雖然能夠通過，但是耗時比較多，其實我們可以再做一個優(yōu)化。
主要看backtrack函數的遞歸部分：
for (int i = 0; i < bucket.length; i++) {
// 剪枝
if (bucket[i] + nums[index] > target) {
continue;
}
if (backtrack(nums, index + 1, bucket, target)) {
return true;
}
}
如果我們讓盡可能多的情況命中剪枝的那個 if 分支，就可以減少遞歸調用的次數，一定程度上減少時間復雜度。
如何盡可能多的命中這個 if 分支呢？要知道我們的index參數是從 0 開始遞增的，也就是遞歸地從 0 開始遍歷nums數組。
如果我們提前對nums數組排序，把大的數字排在前面，那么大的數字會先被分配到bucket中，對于之后的數字，bucket[i] + nums[index]會更大，更容易觸發(fā)剪枝的 if 條件。
所以可以在之前的代碼中再添加一些代碼：
public boolean canPartitionKSubsets(int[] nums, int k) {
// 其他代碼不變
// ...
/* 降序排序 nums 數組 */
Arrays.sort(nums);
int i = 0, j = nums.length - 1;
for (; i < j; i++, j--) {
// 交換 nums[i] 和 nums[j]
int temp = nums[i];
nums[i] = nums[j];
nums[j] = temp;
}
/*******************/
return backtrack(nums, 0, bucket, target);
}
由于 Java 的語言特性，這段代碼通過先升序排序再反轉，達到降序排列的目的。
三、以桶的視角
文章開頭說了，以桶的視角進行窮舉，每個桶需要遍歷nums中的所有數字，決定是否把當前數字裝進桶中；當裝滿一個桶之后，還要裝下一個桶，直到所有桶都裝滿為止。
這個思路可以用下面這段代碼表示出來：
// 裝滿所有桶為止
while (k > 0) {
// 記錄當前桶中的數字之和
int bucket = 0;
for (int i = 0; i < nums.length; i++) {
// 決定是否將 nums[i] 放入當前桶中
bucket += nums[i] or 0;
if (bucket == target) {
// 裝滿了一個桶，裝下一個桶
k--;
break;
}
}
}
那么我們也可以把這個 while 循環(huán)改寫成遞歸函數，不過比剛才略微復雜一些，首先寫一個backtrack遞歸函數出來：
boolean backtrack(int k, int bucket,
int[] nums, int start, boolean[] used, int target);
不要被這么多參數嚇到，我會一個個解釋這些參數。如果你能夠透徹理解本文，也能得心應手地寫出這樣的回溯函數。
這個backtrack函數的參數可以這樣解釋：
現(xiàn)在k號桶正在思考是否應該把nums[start]這個元素裝進來；目前k號桶里面已經裝的數字之和為bucket；used標志某一個元素是否已經被裝到桶中；target是每個桶需要達成的目標和。
根據這個函數定義，可以這樣調用backtrack函數：
public boolean canPartitionKSubsets(int[] nums, int k) {
// 排除一些基本情況
if (k > nums.length) return false;
int sum = 0;
for (int v : nums) sum += v;
if (sum % k != 0) return false;
boolean[] used = new boolean[nums.length];
int target = sum / k;
// k 號桶初始什么都沒裝，從 nums[0] 開始做選擇
return backtrack(k, 0, nums, 0, used, target);
}
實現(xiàn)backtrack函數的邏輯之前，再重復一遍，從桶的視角：
1、需要遍歷nums中所有數字，決定哪些數字需要裝到當前桶中。
2、如果當前桶裝滿了（桶內數字和達到target），則讓下一個桶開始執(zhí)行第 1 步。
下面的代碼就實現(xiàn)了這個邏輯：
boolean backtrack(int k, int bucket,
int[] nums, int start, boolean[] used, int target) {
// base case
if (k == 0) {
// 所有桶都被裝滿了，而且 nums 一定全部用完了
// 因為 target == sum / k
return true;
}
if (bucket == target) {
// 裝滿了當前桶，遞歸窮舉下一個桶的選擇
// 讓下一個桶從 nums[0] 開始選數字
return backtrack(k - 1, 0 ,nums, 0, used, target);
}
// 從 start 開始向后探查有效的 nums[i] 裝入當前桶
for (int i = start; i < nums.length; i++) {
// 剪枝
if (used[i]) {
// nums[i] 已經被裝入別的桶中
continue;
}
if (nums[i] + bucket > target) {
// 當前桶裝不下 nums[i]
continue;
}
// 做選擇，將 nums[i] 裝入當前桶中
used[i] = true;
bucket += nums[i];
// 遞歸窮舉下一個數字是否裝入當前桶
if (backtrack(k, bucket, nums, i + 1, used, target)) {
return true;
}
// 撤銷選擇
used[i] = false;
bucket -= nums[i];
}
// 窮舉了所有數字，都無法裝滿當前桶
return false;
}
至此，這道題的第二種思路也完成了。
四、最后總結
本文寫的這兩種思路都可以通過所有測試用例，不過第一種解法即便經過了排序優(yōu)化，也明顯比第二種解法慢很多，這是為什么呢？
我們來分析一下這兩個算法的時間復雜度，假設nums中的元素個數為n。
先說第一個解法，也就是從數字的角度進行窮舉，n個數字，每個數字有k個桶可供選擇，所以組合出的結果個數為k^n，時間復雜度也就是O(k^n)。
第二個解法，每個桶要遍歷n個數字，選擇「裝入」或「不裝入」，組合的結果有2^n種；而我們有k個桶，所以總的時間復雜度為O(k*2^n)。
當然，這是理論上的最壞復雜度，實際的復雜度肯定要好一些，畢竟我們添加了這么多剪枝邏輯。不過，從復雜度的上界已經可以看出第一種思路要慢很多了。
所以，誰說回溯算法沒有技巧性的？雖然回溯算法就是暴力窮舉，但窮舉也分聰明的窮舉方式和低效的窮舉方式，關鍵看你以誰的「視角」進行窮舉。
通俗來說，我們應該盡量「少量多次」，就是說寧可多做幾次選擇，也不要給太大的選擇空間；寧可「二選一」選k次，也不要 「k選一」選一次。
這道題我們從兩種視角進行窮舉，雖然代碼量看起來多，但核心邏輯都是類似的，相信你通過本文能夠更深刻地理解回溯算法。
推薦閱讀：
專注服務器后臺技術棧知識總結分享
歡迎關注交流共同進步