一道非常經(jīng)典的回溯算法問題：子集劃分問題

遞歸遍歷數(shù)組你會不會？其實也很簡單：

voidtraverse(int[]nums,intindex){
if(index==nums.length){
return;
}
System.out.println(nums[index]);
traverse(nums,index+1);
}

只要調(diào)用traverse(nums, 0)，和 for 循環(huán)的效果是完全一樣的。

那么回到這道題，以數(shù)字的視角，選擇k個桶，用 for 循環(huán)寫出來是下面這樣：

//k個桶（集合），記錄每個桶裝的數(shù)字之和
int[]bucket=newint[k];

//窮舉nums中的每個數(shù)字
for(intindex=0;index//窮舉每個桶
for(inti=0;i//nums[index]選擇是否要進(jìn)入第i個桶
//...
}
}

如果改成遞歸的形式，就是下面這段代碼邏輯：

//k個桶（集合），記錄每個桶裝的數(shù)字之和
int[]bucket=newint[k];

//窮舉nums中的每個數(shù)字
voidbacktrack(int[]nums,intindex){
//basecase
if(index==nums.length){
return;
}
//窮舉每個桶
for(inti=0;i//選擇裝進(jìn)第i個桶
bucket[i]+=nums[index];
//遞歸窮舉下一個數(shù)字的選擇
backtrack(nums,index+1);
//撤銷選擇
bucket[i]-=nums[index];
}
}

雖然上述代碼僅僅是窮舉邏輯，還不能解決我們的問題，但是只要略加完善即可：

//主函數(shù)
booleancanPartitionKSubsets(int[]nums,intk){
//排除一些基本情況
if(k>nums.length)returnfalse;
intsum=0;
for(intv:nums)sum+=v;
if(sum%k!=0)returnfalse;

//k個桶（集合），記錄每個桶裝的數(shù)字之和
int[]bucket=newint[k];
//理論上每個桶（集合）中數(shù)字的和
inttarget=sum/k;
//窮舉，看看nums是否能劃分成k個和為target的子集
returnbacktrack(nums,0,bucket,target);
}

//遞歸窮舉nums中的每個數(shù)字
booleanbacktrack(
int[]nums,intindex,int[]bucket,inttarget){

if(index==nums.length){
//檢查所有桶的數(shù)字之和是否都是target
for(inti=0;iif(bucket[i]!=target){
returnfalse;
}
}
//nums成功平分成k個子集
returntrue;
}

//窮舉nums[index]可能裝入的桶
for(inti=0;i//剪枝，桶裝裝滿了
if(bucket[i]+nums[index]>target){
continue;
}
//將nums[index]裝入bucket[i]
bucket[i]+=nums[index];
//遞歸窮舉下一個數(shù)字的選擇
if(backtrack(nums,index+1,bucket,target)){
returntrue;
}
//撤銷選擇
bucket[i]-=nums[index];
}

//nums[index]裝入哪個桶都不行
returnfalse;
}

有之前的鋪墊，相信這段代碼是比較容易理解的。這個解法雖然能夠通過，但是耗時比較多，其實我們可以再做一個優(yōu)化。

主要看backtrack函數(shù)的遞歸部分：

for(inti=0;i//剪枝
if(bucket[i]+nums[index]>target){
continue;
}

if(backtrack(nums,index+1,bucket,target)){
returntrue;
}
}

如果我們讓盡可能多的情況命中剪枝的那個 if 分支，就可以減少遞歸調(diào)用的次數(shù)，一定程度上減少時間復(fù)雜度。

如何盡可能多的命中這個 if 分支呢？要知道我們的index參數(shù)是從 0 開始遞增的，也就是遞歸地從 0 開始遍歷nums數(shù)組。

如果我們提前對nums數(shù)組排序，把大的數(shù)字排在前面，那么大的數(shù)字會先被分配到bucket中，對于之后的數(shù)字，bucket[i] + nums[index]會更大，更容易觸發(fā)剪枝的 if 條件。

所以可以在之前的代碼中再添加一些代碼：

booleancanPartitionKSubsets(int[]nums,intk){
//其他代碼不變
//...
/*降序排序nums數(shù)組*/
Arrays.sort(nums);
for(i=0,j=nums.length-1;i//交換nums[i]和nums[j]
inttemp=nums[i];
nums[i]=nums[j];
nums[j]=temp;
}
/*******************/
returnbacktrack(nums,0,bucket,target);
}

由于 Java 的語言特性，這段代碼通過先升序排序再反轉(zhuǎn)，達(dá)到降序排列的目的。

三、以桶的視角

文章開頭說了，以桶的視角進(jìn)行窮舉，每個桶需要遍歷nums中的所有數(shù)字，決定是否把當(dāng)前數(shù)字裝進(jìn)桶中；當(dāng)裝滿一個桶之后，還要裝下一個桶，直到所有桶都裝滿為止。

這個思路可以用下面這段代碼表示出來：

//裝滿所有桶為止
while(k>0){
//記錄當(dāng)前桶中的數(shù)字之和
intbucket=0;
for(inti=0;i//決定是否將nums[i]放入當(dāng)前桶中
bucket+=nums[i]or0;
if(bucket==target){
//裝滿了一個桶，裝下一個桶
k--;
break;
}
}
}

那么我們也可以把這個 while 循環(huán)改寫成遞歸函數(shù)，不過比剛才略微復(fù)雜一些，首先寫一個backtrack遞歸函數(shù)出來：

booleanbacktrack(intk,intbucket,
int[]nums,intstart,boolean[]used,inttarget);

不要被這么多參數(shù)嚇到，我會一個個解釋這些參數(shù)。如果你能夠透徹理解本文，也能得心應(yīng)手地寫出這樣的回溯函數(shù)。

這個backtrack函數(shù)的參數(shù)可以這樣解釋：

現(xiàn)在k號桶正在思考是否應(yīng)該把nums[start]這個元素裝進(jìn)來；目前k號桶里面已經(jīng)裝的數(shù)字之和為bucket；used標(biāo)志某一個元素是否已經(jīng)被裝到桶中；target是每個桶需要達(dá)成的目標(biāo)和。

根據(jù)這個函數(shù)定義，可以這樣調(diào)用backtrack函數(shù)：

booleancanPartitionKSubsets(int[]nums,intk){
//排除一些基本情況
if(k>nums.length)returnfalse;
intsum=0;
for(intv:nums)sum+=v;
if(sum%k!=0)returnfalse;

boolean[]used=newboolean[nums.length];
inttarget=sum/k;
//k號桶初始什么都沒裝，從nums[0]開始做選擇
returnbacktrack(k,0,nums,0,used,target);
}

實現(xiàn)backtrack函數(shù)的邏輯之前，再重復(fù)一遍，從桶的視角：

1、需要遍歷nums中所有數(shù)字，決定哪些數(shù)字需要裝到當(dāng)前桶中。

2、如果當(dāng)前桶裝滿了（桶內(nèi)數(shù)字和達(dá)到target），則讓下一個桶開始執(zhí)行第 1 步。

下面的代碼就實現(xiàn)了這個邏輯：

booleanbacktrack(intk,intbucket,
int[]nums,intstart,boolean[]used,inttarget){
//basecase
if(k==0){
//所有桶都被裝滿了，而且nums一定全部用完了
//因為target==sum/k
returntrue;
}
if(bucket==target){
//裝滿了當(dāng)前桶，遞歸窮舉下一個桶的選擇
//讓下一個桶從nums[0]開始選數(shù)字
returnbacktrack(k-1,0,nums,0,used,target);
}

//從start開始向后探查有效的nums[i]裝入當(dāng)前桶
for(inti=start;i//剪枝
if(used[i]){
//nums[i]已經(jīng)被裝入別的桶中
continue;
}
if(nums[i]+bucket>target){
//當(dāng)前桶裝不下nums[i]
continue;
}
//做選擇，將nums[i]裝入當(dāng)前桶中
used[i]=true;
bucket+=nums[i];
//遞歸窮舉下一個數(shù)字是否裝入當(dāng)前桶
if(backtrack(k,bucket,nums,i+1,used,target)){
returntrue;
}
//撤銷選擇
used[i]=false;
bucket-=nums[i];
}
//窮舉了所有數(shù)字，都無法裝滿當(dāng)前桶
returnfalse;
}

這段代碼是可以得出正確答案的，但是效率很低，我們可以思考一下是否還有優(yōu)化的空間。

首先，在這個解法中每個桶都可以認(rèn)為是沒有差異的，但是我們的回溯算法卻會對它們區(qū)別對待，這里就會出現(xiàn)重復(fù)計算的情況。

什么意思呢？我們的回溯算法，說到底就是窮舉所有可能的組合，然后看是否能找出和為target的k個桶（子集）。

那么，比如下面這種情況，target = 5，算法會在第一個桶里面裝1, 4：

現(xiàn)在第一個桶裝滿了，就開始裝第二個桶，算法會裝入2, 3：

然后以此類推，對后面的元素進(jìn)行窮舉，湊出若干個和為 5 的桶（子集）。

但問題是，如果最后發(fā)現(xiàn)無法湊出和為target的k個子集，算法會怎么做？

回溯算法會回溯到第一個桶，重新開始窮舉，現(xiàn)在它知道第一個桶里裝1, 4是不可行的，它會嘗試把2, 3裝到第一個桶里：

現(xiàn)在第一個桶裝滿了，就開始裝第二個桶，算法會裝入1, 4：

好，到這里你應(yīng)該看出來問題了，這種情況其實和之前的那種情況是一樣的。也就是說，到這里你其實已經(jīng)知道不需要再窮舉了，必然湊不出來和為target的k個子集。

但我們的算法還是會傻乎乎地繼續(xù)窮舉，因為在她看來，第一個桶和第二個桶里面裝的元素不一樣，那這就是兩種不一樣的情況呀。

那么我們怎么讓算法的智商提高，識別出這種情況，避免冗余計算呢？

你注意這兩種情況的used數(shù)組肯定長得一樣，所以used數(shù)組可以認(rèn)為是回溯過程中的「狀態(tài)」。

所以，我們可以用一個memo備忘錄，在裝滿一個桶時記錄當(dāng)前used的狀態(tài)，如果當(dāng)前used的狀態(tài)是曾經(jīng)出現(xiàn)過的，那就不用再繼續(xù)窮舉，從而起到剪枝避免冗余計算的作用。

有讀者肯定會問，used是一個布爾數(shù)組，怎么作為鍵進(jìn)行存儲呢？這其實是小問題，比如我們可以把數(shù)組轉(zhuǎn)化成字符串，這樣就可以作為哈希表的鍵進(jìn)行存儲了。

看下代碼實現(xiàn)，只要稍微改一下backtrack函數(shù)即可：

//備忘錄，存儲used數(shù)組的狀態(tài)
HashMapmemo=newHashMap<>();

booleanbacktrack(intk,intbucket,int[]nums,intstart,boolean[]used,inttarget){
//basecase
if(k==0){
returntrue;
}
//將used的狀態(tài)轉(zhuǎn)化成形如[true,false,...]的字符串
//便于存入HashMap
Stringstate=Arrays.toString(used);

if(bucket==target){
//裝滿了當(dāng)前桶，遞歸窮舉下一個桶的選擇
booleanres=backtrack(k-1,0,nums,0,used,target);
//將當(dāng)前狀態(tài)和結(jié)果存入備忘錄
memo.put(state,res);
returnres;
}

if(memo.containsKey(state)){
//如果當(dāng)前狀態(tài)曾今計算過，就直接返回，不要再遞歸窮舉了
returnmemo.get(state);
}

//其他邏輯不變...
}

這樣提交解法，發(fā)現(xiàn)執(zhí)行效率依然比較低，這次不是因為算法邏輯上的冗余計算，而是代碼實現(xiàn)上的問題。

因為每次遞歸都要把used數(shù)組轉(zhuǎn)化成字符串，這對于編程語言來說也是一個不小的消耗，所以我們還可以進(jìn)一步優(yōu)化。

注意題目給的數(shù)據(jù)規(guī)模nums.length <= 16，也就是說used數(shù)組最多也不會超過 16，那么我們完全可以用「位圖」的技巧，用一個 int 類型的used變量來替代used數(shù)組。

具體來說，我們可以用整數(shù)used的第i位（(used >> i) & 1）的 1/0 來表示used[i]的 true/false。

這樣一來，不僅節(jié)約了空間，而且整數(shù)used也可以直接作為鍵存入 HashMap，省去數(shù)組轉(zhuǎn)字符串的消耗。

看下最終的解法代碼：

publicbooleancanPartitionKSubsets(int[]nums,intk){
//排除一些基本情況
if(k>nums.length)returnfalse;
intsum=0;
for(intv:nums)sum+=v;
if(sum%k!=0)returnfalse;

intused=0;//使用位圖技巧
inttarget=sum/k;
//k號桶初始什么都沒裝，從nums[0]開始做選擇
returnbacktrack(k,0,nums,0,used,target);
}

HashMapmemo=newHashMap<>();

booleanbacktrack(intk,intbucket,
int[]nums,intstart,intused,inttarget){
//basecase
if(k==0){
//所有桶都被裝滿了，而且nums一定全部用完了
returntrue;
}
if(bucket==target){
//裝滿了當(dāng)前桶，遞歸窮舉下一個桶的選擇
//讓下一個桶從nums[0]開始選數(shù)字
booleanres=backtrack(k-1,0,nums,0,used,target);
//緩存結(jié)果
memo.put(used,res);
returnres;
}

if(memo.containsKey(used)){
//避免冗余計算
returnmemo.get(used);
}

for(inti=start;i//剪枝
if(((used>>i)&1)==1){//判斷第i位是否是1
//nums[i]已經(jīng)被裝入別的桶中
continue;
}
if(nums[i]+bucket>target){
continue;
}
//做選擇
used|=1<//將第i位置為1
bucket+=nums[i];
//遞歸窮舉下一個數(shù)字是否裝入當(dāng)前桶
if(backtrack(k,bucket,nums,i+1,used,target)){
returntrue;
}
//撤銷選擇
used^=1<//使用異或運(yùn)算將第i位恢復(fù)0
bucket-=nums[i];
}

returnfalse;
}

至此，這道題的第二種思路也完成了。

四、最后總結(jié)

本文寫的這兩種思路都可以算出正確答案，不過第一種解法即便經(jīng)過了排序優(yōu)化，也明顯比第二種解法慢很多，這是為什么呢？

我們來分析一下這兩個算法的時間復(fù)雜度，假設(shè)nums中的元素個數(shù)為n。

先說第一個解法，也就是從數(shù)字的角度進(jìn)行窮舉，n個數(shù)字，每個數(shù)字有k個桶可供選擇，所以組合出的結(jié)果個數(shù)為k^n，時間復(fù)雜度也就是O(k^n)。

第二個解法，每個桶要遍歷n個數(shù)字，對每個數(shù)字有「裝入」或「不裝入」兩種選擇，所以組合的結(jié)果有2^n種；而我們有k個桶，所以總的時間復(fù)雜度為O(k*2^n)。

當(dāng)然，這是對最壞復(fù)雜度上界的粗略估算，實際的復(fù)雜度肯定要好很多，畢竟我們添加了這么多剪枝邏輯。不過，從復(fù)雜度的上界已經(jīng)可以看出第一種思路要慢很多了。

所以，誰說回溯算法沒有技巧性的？雖然回溯算法就是暴力窮舉，但窮舉也分聰明的窮舉方式和低效的窮舉方式，關(guān)鍵看你以誰的「視角」進(jìn)行窮舉。

通俗來說，我們應(yīng)該盡量「少量多次」，就是說寧可多做幾次選擇，也不要給太大的選擇空間；寧可「二選一」選k次，也不要 「k選一」選一次。

好了，這道題我們從兩種視角進(jìn)行窮舉，雖然代碼量看起來多，但核心邏輯都是類似的，相信你通過本文能夠更深刻地理解回溯算法。