1. 從一個(gè)“簡單”的計(jì)數(shù)問題說起最近在帶新人刷算法題遇到一個(gè)經(jīng)典問題它看起來人畜無害卻讓不少初學(xué)者栽了跟頭。題目大意是這樣的給你一個(gè)長度為n的格子你需要用k種顏色去涂滿它。但有一個(gè)限制相鄰的兩個(gè)格子不能涂成相同的顏色。問一共有多少種不同的涂色方案乍一看這不就是個(gè)排列組合題嗎第一個(gè)格子有k種選擇第二個(gè)格子不能和第一個(gè)相同所以有k-1種選擇以此類推。總方案數(shù)不就是k * (k-1)^(n-1)嗎這個(gè)公式在n和k都不大的時(shí)候確實(shí)能快速給出答案。很多新人做到這里就心滿意足地提交了然后……就收到了一個(gè)“Wrong Answer”。問題出在哪這個(gè)公式成立的前提是顏色是“無限”的或者說我們每次選擇時(shí)可用的顏色數(shù)量只受“上一個(gè)格子顏色”這一個(gè)條件的約束。但在很多實(shí)際問題中約束條件要復(fù)雜得多。比如如果顏色數(shù)量k很小或者題目增加了額外的限制比如“首尾格子也不能同色”甚至“某些特定位置的格子有固定顏色要求”剛才那個(gè)簡單的乘法原理就立刻失效了。這時(shí)我們面對的就不再是一個(gè)有閉合公式的問題而是一個(gè)需要系統(tǒng)化搜索所有可能狀態(tài)的問題。這就是“著色方案”類問題的核心在滿足一系列復(fù)雜約束條件下計(jì)算所有可行的分配方案總數(shù)。當(dāng)約束變得具體而微暴力枚舉所有可能性在數(shù)據(jù)規(guī)模稍大時(shí)就會變得不可能時(shí)間復(fù)雜度是k^n的指數(shù)級。此時(shí)我們亟需一種更聰明的方法而“記憶化搜索”正是為此而生的利器。它不是什么高深莫測的黑魔法而是我們面對復(fù)雜狀態(tài)空間時(shí)一種化繁為簡、避免重復(fù)勞動的樸素思想。接下來我們就剝開這層外衣看看它到底是怎么工作的以及如何用它來優(yōu)雅地解決那些看似棘手的計(jì)數(shù)問題。2. 暴力搜索的困境與狀態(tài)定義的藝術(shù)在討論記憶化搜索之前我們必須先理解它所試圖優(yōu)化的對象——深度優(yōu)先搜索DFS。對于著色問題最直接的思路就是遞歸回溯從第一個(gè)格子開始嘗試每一種可能的顏色如果當(dāng)前選擇不違反約束比如和左邊格子顏色不同就遞歸地去涂下一個(gè)格子。當(dāng)所有格子都涂滿時(shí)就得到了一種合法方案計(jì)數(shù)器加一。def dfs(position, n, k, prev_color): if position n: # 所有格子涂完找到一種方案 return 1 total 0 for color in range(k): if color ! prev_color: # 簡單相鄰約束 total dfs(position 1, n, k, color) return total # 初始調(diào)用假設(shè)第一個(gè)格子左邊沒有格子prev_color用-1表示 result dfs(0, n, k, -1)這段代碼清晰易懂但它有一個(gè)致命缺陷存在大量重復(fù)計(jì)算。舉個(gè)例子假設(shè)n5, k3。當(dāng)我們遞歸探索時(shí)可能會先走顏色序列A-B-?這條路徑計(jì)算完后面所有的可能性。之后在另一條分支里我們可能又遇到了顏色序列C-B-?的狀態(tài)。注意此時(shí)雖然前兩個(gè)格子的顏色不同A和C但第二個(gè)格子都是B并且我們即將面對的是第三個(gè)格子。對于從“第三個(gè)格子開始前一個(gè)顏色是B”這個(gè)子問題它的答案是完全一樣的與第一個(gè)格子是A還是C無關(guān)然而我們的樸素DFS會傻乎乎地重新計(jì)算一遍。這就是狀態(tài)重疊。我們遞歸函數(shù)的本質(zhì)是在計(jì)算一個(gè)(當(dāng)前位置, 前一個(gè)格子顏色)所確定的子問題的解。一旦這個(gè)二元組(pos, prev_color)確定了無論通過哪條路徑到達(dá)這個(gè)狀態(tài)后續(xù)的涂色方案數(shù)都是唯一確定的。如果我們能把這個(gè)結(jié)果存起來下次再遇到相同的(pos, prev_color)時(shí)直接返回結(jié)果就能節(jié)省巨大的計(jì)算量。所以記憶化搜索的第一步也是最重要的一步就是精確定義“狀態(tài)”。狀態(tài)必須能唯一標(biāo)識一個(gè)子問題并且其數(shù)量是可控的。對于基本的相鄰不同色問題狀態(tài)就是(pos, prev_color)。其中pos的范圍是0到nn表示已涂完是遞歸終點(diǎn)prev_color的范圍是k種顏色再加上一個(gè)表示“無前驅(qū)”的特殊值比如-1。因此狀態(tài)總數(shù)大約是(n1) * (k1)這是一個(gè)多項(xiàng)式級別遠(yuǎn)遠(yuǎn)小于指數(shù)級的k^n。注意狀態(tài)定義并非一成不變。如果約束變成“首尾不能同色”我們的狀態(tài)就需要增加信息比如變成(pos, prev_color, first_color)因?yàn)樽詈笠粋€(gè)格子的選擇受第一個(gè)格子顏色的影響。定義狀態(tài)的關(guān)鍵在于找出哪些信息是決定后續(xù)選擇所必需的、最小的信息集合。這需要根據(jù)具體問題的約束條件進(jìn)行設(shè)計(jì)和提煉是記憶化搜索中最具技巧性的部分。3. 記憶化搜索的實(shí)現(xiàn)框架與細(xì)節(jié)打磨理解了狀態(tài)實(shí)現(xiàn)記憶化搜索就水到渠成了。我們用一個(gè)緩存通常是一個(gè)字典或數(shù)組來存儲已經(jīng)計(jì)算過的狀態(tài)結(jié)果。這個(gè)緩存結(jié)構(gòu)的選擇很有講究。1. 緩存數(shù)據(jù)結(jié)構(gòu)的選擇字典Dict/HashMap最通用和靈活。鍵Key是狀態(tài)值Value是結(jié)果。當(dāng)狀態(tài)比較復(fù)雜比如包含多個(gè)離散變量時(shí)用字典很自然。例如狀態(tài)(pos, prev_color)可以轉(zhuǎn)化為元組(pos, prev_color)作為鍵。多維數(shù)組List/Array當(dāng)狀態(tài)的所有維度都是整數(shù)且范圍明確時(shí)使用數(shù)組訪問效率更高。例如pos范圍[0, n]prev_color范圍[-1, k-1]我們可以建立一個(gè)(n1) x (k1)的二維數(shù)組dp其中dp[pos][prev_color1]存儲結(jié)果1是為了將-1映射到索引0。在著色方案這類典型問題中狀態(tài)維度固定且范圍小使用數(shù)組是更優(yōu)解。它不僅速度快而且代碼清晰。2. 遞歸函數(shù)的改造我們將樸素的DFS函數(shù)改造成一個(gè)“有記憶”的DFS。第一步查緩存。在函數(shù)開始時(shí)先檢查當(dāng)前狀態(tài)是否已經(jīng)計(jì)算過。如果是直接返回緩存的結(jié)果。第二步遞歸計(jì)算。如果沒計(jì)算過則進(jìn)行正常的遞歸邏輯計(jì)算所有可能的選擇并求和。第三步存緩存。在返回結(jié)果之前將(當(dāng)前狀態(tài), 計(jì)算結(jié)果)存入緩存。以下是使用二維數(shù)組作為緩存的經(jīng)典實(shí)現(xiàn)def count_colorings(n, k): # dp[pos][prev_color1], 初始化所有值為-1表示未計(jì)算 # prev_color 從 -1 到 k-1所以第二維大小是 k1 dp [[-1] * (k 1) for _ in range(n 1)] def dfs(pos, prev_color_idx): # prev_color_idx 是 prev_color 在dp數(shù)組中的索引 (prev_color 1) if pos n: return 1 # 成功涂完所有格子找到一種方案 if dp[pos][prev_color_idx] ! -1: return dp[pos][prev_color_idx] total 0 for color in range(k): # 將顏色值color轉(zhuǎn)換為“前一個(gè)顏色”的索引表示用于比較 # 注意prev_color_idx 是索引真正的 prev_color prev_color_idx - 1 actual_prev_color prev_color_idx - 1 if color ! actual_prev_color: # 遞歸下一個(gè)位置的前一個(gè)顏色索引是 color 1 total dfs(pos 1, color 1) dp[pos][prev_color_idx] total return total # 初始調(diào)用從位置0開始前一個(gè)顏色不存在用索引0表示即 actual_prev_color -1 return dfs(0, 0) # 示例5個(gè)格子3種顏色相鄰不同色 print(count_colorings(5, 3)) # 輸出應(yīng)為 3 * 2^4 48可以用公式驗(yàn)證3. 邊界條件與初始化遞歸的終點(diǎn)pos n通常返回 1表示找到一種完整方案。緩存數(shù)組的初始化值必須是一個(gè)不會出現(xiàn)在正常結(jié)果中的值如-1用以區(qū)分“未計(jì)算”和“計(jì)算結(jié)果為0”后者在某些問題中是合法結(jié)果表示無解。4. 復(fù)雜度分析時(shí)間復(fù)雜度由于每個(gè)狀態(tài)(pos, prev_color)最多只計(jì)算一次每次計(jì)算需要遍歷k種顏色所以總時(shí)間復(fù)雜度為O(n * k * k)等等仔細(xì)看內(nèi)層循環(huán)。對于每個(gè)狀態(tài)我們循環(huán)k次每次遞歸調(diào)用是 O(1) 的查表或計(jì)算。因此準(zhǔn)確的時(shí)間復(fù)雜度是O(狀態(tài)數(shù) * 每個(gè)狀態(tài)的計(jì)算成本) O(n * k * 1) O(n * k)。這里的k是顏色數(shù)通常是個(gè)常數(shù)或者不大的數(shù)因此算法是線性或近似線性的效率極高?？臻g復(fù)雜度主要是緩存數(shù)組dp的開銷為O(n * k)以及遞歸調(diào)用棧的深度O(n)。實(shí)操心得在實(shí)現(xiàn)時(shí)我強(qiáng)烈建議將“狀態(tài)”到“緩存索引”的映射關(guān)系單獨(dú)寫成一個(gè)清晰的函數(shù)或注釋。比如get_index(prev_color)。這能極大減少因?yàn)橄聵?biāo)轉(zhuǎn)換錯(cuò)誤導(dǎo)致的Bug尤其是在狀態(tài)變量有特殊值如-1的時(shí)候。另外對于結(jié)果可能非常大的計(jì)數(shù)問題比如方案數(shù)可能超過64位整數(shù)范圍要在題目要求下及時(shí)取模并且在存入緩存和返回結(jié)果前都要取模保證一致性。4. 從經(jīng)典到變種應(yīng)對更復(fù)雜的約束條件記憶化搜索的強(qiáng)大之處在于其靈活性。當(dāng)問題的約束條件發(fā)生變化時(shí)我們通常不需要推翻重來而只需調(diào)整“狀態(tài)定義”和“狀態(tài)轉(zhuǎn)移”邏輯。下面我們通過幾個(gè)變種問題來體會這一點(diǎn)。4.1 變種一首尾格子也不能同色這是“相鄰不同色”問題的經(jīng)典加強(qiáng)版。此時(shí)最后一個(gè)格子第n-1個(gè)的顏色不僅不能和它左邊的格子第n-2個(gè)相同還不能和第一個(gè)格子相同。狀態(tài)定義的升級原來的狀態(tài)(pos, prev_color)不足以決定最后一個(gè)格子的選擇因?yàn)樗鄙倭恕暗谝粋€(gè)格子顏色”的信息。因此我們需要將“第一個(gè)格子的顏色”也納入狀態(tài)。定義狀態(tài)為(pos, prev_color, first_color)。其中first_color在遞歸開始時(shí)就確定下來并一路傳遞下去。狀態(tài)轉(zhuǎn)移的調(diào)整在遞歸涂色時(shí)當(dāng)pos 0涂第一個(gè)格子遍歷所有k種顏色作為first_color同時(shí)這個(gè)顏色也是prev_color。當(dāng)pos n-1涂最后一個(gè)格子遍歷顏色時(shí)除了要滿足color ! prev_color還必須滿足color ! first_color。其他位置和之前一樣只需滿足color ! prev_color。緩存維度狀態(tài)變成了三維(pos, prev_color, first_color)緩存數(shù)組的大小變?yōu)?n) * (k) * (k)。雖然空間變大了但相對于指數(shù)爆炸這依然是完全可以接受的。4.2 變種二顏色使用次數(shù)限制假設(shè)每種顏色最多只能使用m次。這在實(shí)際場景中很常見比如有限的顏料庫存。狀態(tài)定義的升級此時(shí)僅僅知道前一個(gè)顏色是什么不夠了我們還需要知道每種顏色還剩多少使用次數(shù)。一種直觀的狀態(tài)定義是(pos, prev_color, color_used_tuple)其中color_used_tuple是一個(gè)長度為k的元組記錄每種顏色已使用的次數(shù)。但這樣狀態(tài)空間會非常大n * k * (m1)^k。優(yōu)化思路對于計(jì)數(shù)問題我們往往不需要知道每種顏色具體用了多少次而只需要知道“剩余使用次數(shù)”的模式。如果所有顏色的限制次數(shù)m相同那么問題可以簡化為在涂到某個(gè)位置時(shí)有多少種顏色已經(jīng)用滿了m次有多少種顏色用了m-1次……但這依然復(fù)雜。一個(gè)更實(shí)用的方法是當(dāng)k和m不大時(shí)可以使用狀態(tài)壓縮。用一個(gè)k位的整數(shù)比特位來表示哪些顏色已經(jīng)用盡了次數(shù)或者用一個(gè)整數(shù)數(shù)組來記錄使用次數(shù)并將整個(gè)數(shù)組作為字典的鍵雖然效率會降低。這體現(xiàn)了記憶化搜索的另一個(gè)維度當(dāng)狀態(tài)本身復(fù)雜時(shí)我們可以利用哈希表字典的靈活性來存儲。4.3 變種三格子分組著色圖著色問題的簡化問題升級為格子之間不是簡單的線性關(guān)系而是一個(gè)一般的圖。每個(gè)節(jié)點(diǎn)格子需要著色有邊相連的節(jié)點(diǎn)不能同色。這就是經(jīng)典的圖著色問題是NP難的。但對于特定的、樹狀或稀疏的圖記憶化搜索結(jié)合樹形DP仍然可以高效解決。狀態(tài)定義對于樹形結(jié)構(gòu)我們通常在樹上進(jìn)行DFS。狀態(tài)可以定義為(node, parent_color)表示在以node為根的子樹中當(dāng)node的父節(jié)點(diǎn)顏色為parent_color時(shí)該子樹的著色方案數(shù)。然后通過遞歸合并子節(jié)點(diǎn)的結(jié)果來計(jì)算當(dāng)前節(jié)點(diǎn)的方案數(shù)。踩坑實(shí)錄在處理復(fù)雜約束時(shí)最容易犯的錯(cuò)誤是狀態(tài)定義遺漏了關(guān)鍵信息。我曾在一個(gè)比賽中遇到一個(gè)問題要求“任意兩個(gè)距離為2的格子也不能同色”。我最初只定義了(pos, prev_color)結(jié)果總是少算。后來才意識到距離為2意味著當(dāng)前格子不能和它前面第2個(gè)格子同色。因此狀態(tài)必須包含前兩個(gè)格子的顏色信息即(pos, color_of_pos_minus_1, color_of_pos_minus_2)。這個(gè)教訓(xùn)讓我明白定義狀態(tài)時(shí)要像偵探一樣問自己“要唯一確定從現(xiàn)在開始的所有未來可能性最少需要知道過去的哪些信息”5. 記憶化搜索 vs. 動態(tài)規(guī)劃思維路徑的異同很多人會把記憶化搜索和動態(tài)規(guī)劃DP等同起來稱其為“遞歸形式的DP”。這種說法有一定道理但兩者在思維起點(diǎn)和實(shí)現(xiàn)方式上有著微妙的區(qū)別理解這些區(qū)別能幫助你更好地選擇工具。5.1 思維路徑的對比記憶化搜索Memoization思維是自頂向下的。你從要解決的原問題如f(0, -1)開始思考“要解決我的問題我需要先解決哪些子問題”然后遞歸地去解決這些子問題并用緩存避免重復(fù)。它的思路更符合人類面對復(fù)雜問題的自然分解過程——分而治之。動態(tài)規(guī)劃Dynamic Programming思維是自底向上的。你需要先確定所有子問題的計(jì)算順序通常是較小的、基礎(chǔ)的狀態(tài)先計(jì)算然后通過迭代循環(huán)從小問題逐步推導(dǎo)出大問題的解。這需要更強(qiáng)的“全局”狀態(tài)轉(zhuǎn)移視角。對于著色方案問題記憶化搜索的思維是“我想知道從第0個(gè)格子開始涂有幾種方案。那我先試試涂第一種顏色然后問題就變成了‘從第1個(gè)格子開始且前一個(gè)顏色是第一種顏色’有幾種方案。我去計(jì)算這個(gè)子問題……”而動態(tài)規(guī)劃則會先計(jì)算“最后一個(gè)格子怎么涂”然后倒推回來或者從第一個(gè)格子開始正推。5.2 實(shí)現(xiàn)形式的對比我們以基礎(chǔ)著色問題為例看看兩者的代碼實(shí)現(xiàn)。記憶化搜索遞歸如上文所示代碼直觀反映了遞歸關(guān)系。動態(tài)規(guī)劃迭代我們需要定義dp[i][c]表示“涂完前i個(gè)格子并且第i個(gè)格子最后一個(gè)顏色是c的方案總數(shù)”。狀態(tài)轉(zhuǎn)移dp[i][c] sum(dp[i-1][c])其中c是所有不等于c的顏色。因?yàn)榈趇個(gè)格子涂c那么第i-1個(gè)格子可以是任何非c的顏色。初始化dp[0][c] 1對于第一個(gè)格子每種顏色都是一種方案。最終答案sum(dp[n-1][c])對所有顏色c求和。def count_colorings_dp(n, k): if n 0: return 0 # dp[i][c]: 前i個(gè)格子已涂完且第i個(gè)格子顏色為c的方案數(shù) (i從0開始) dp [[0] * k for _ in range(n)] # 初始化第一個(gè)格子 for c in range(k): dp[0][c] 1 # 遞推 for i in range(1, n): for c in range(k): # 當(dāng)前格子涂c上一個(gè)格子可以涂任何非c的顏色 for prev_c in range(k): if prev_c ! c: dp[i][c] dp[i-1][prev_c] # 總和 total sum(dp[n-1][c] for c in range(k)) return total5.3 如何選擇優(yōu)先考慮記憶化搜索當(dāng)狀態(tài)轉(zhuǎn)移關(guān)系不那么直觀或者存在復(fù)雜的依賴關(guān)系比如在樹上記憶化搜索更容易思考和實(shí)現(xiàn)。你只需要寫出遞歸關(guān)系讓緩存去處理重復(fù)子問題?？紤]動態(tài)規(guī)劃當(dāng)問題有明顯的線性順序且狀態(tài)轉(zhuǎn)移方程清晰簡單時(shí)DP的迭代形式通常效率稍高避免了遞歸調(diào)用開銷并且不容易出現(xiàn)棧溢出對于深度很大的遞歸。一個(gè)實(shí)用的建議先嘗試用記憶化搜索的思路去思考和解決問題。寫出遞歸函數(shù)。如果發(fā)現(xiàn)性能或棧深度有問題再考慮是否能轉(zhuǎn)化為等價(jià)的、自底向上的動態(tài)規(guī)劃。很多時(shí)候記憶化搜索是探索DP狀態(tài)轉(zhuǎn)移方程的絕佳跳板。個(gè)人經(jīng)驗(yàn)在競賽或面試中如果時(shí)間緊迫我通常會首選記憶化搜索。因?yàn)樗蝗菀壮鲥e(cuò)思維負(fù)擔(dān)小。只要確保狀態(tài)定義正確、緩存生效基本就能拿到分?jǐn)?shù)。而自底向上的DP一旦遞推順序或初始化寫錯(cuò)調(diào)試起來可能更費(fèi)時(shí)間。當(dāng)然對于狀態(tài)空間巨大、需要滾動數(shù)組優(yōu)化空間的情況就必須使用迭代DP了。6. 性能優(yōu)化與邊界陷阱即使使用了記憶化搜索如果不注意細(xì)節(jié)依然可能掉入性能或正確性的陷阱。這里分享幾個(gè)關(guān)鍵點(diǎn)。6.1 緩存鍵的設(shè)計(jì)與哈希效率當(dāng)使用字典Python的dict或functools.lru_cache時(shí)狀態(tài)的哈希效率至關(guān)重要。最常用的方法是將狀態(tài)轉(zhuǎn)換為元組Tuple。但要注意如果狀態(tài)中包含列表List必須先轉(zhuǎn)換為元組因?yàn)榱斜硎遣豢晒５?。對于整?shù)狀態(tài)直接使用元組即可。對于復(fù)雜對象可以考慮使用字符串編碼或自定義哈希函數(shù)。在Python中使用lru_cache(maxsizeNone)裝飾器可以極簡地實(shí)現(xiàn)記憶化它自動將函數(shù)參數(shù)作為緩存鍵。這對于原型設(shè)計(jì)和快速驗(yàn)證非常方便。from functools import lru_cache lru_cache(maxsizeNone) def dfs(pos, prev_color): if pos n: return 1 total 0 for color in range(k): if color ! prev_color: total dfs(pos 1, color) return total6.2 遞歸深度限制Python默認(rèn)的遞歸深度限制通常為1000對于n較大的問題可能不夠。對于線性遞歸深度為n的問題當(dāng)n超過1000時(shí)需要手動設(shè)置遞歸深度或改用迭代DP。import sys sys.setrecursionlimit(10000) # 設(shè)置為一個(gè)更大的值但更根本的解決方法是評估問題是否必須深度遞歸。像著色方案這種問題遞歸深度等于格子數(shù)n如果n達(dá)到10^5級別即使解除限制遞歸調(diào)用棧的開銷也極大且有棧溢出風(fēng)險(xiǎn)。此時(shí)必須使用迭代的動態(tài)規(guī)劃。6.3 大數(shù)取模的處理方案數(shù)往往非常巨大題目通常要求對某個(gè)大數(shù)MOD如10^97取模。這里有一個(gè)極易出錯(cuò)的細(xì)節(jié)必須在每一次加法運(yùn)算后立即取模而不是最后才取模。因?yàn)橹虚g結(jié)果可能已經(jīng)溢出即使在Python這種大整數(shù)語言中取模操作本身也應(yīng)在合理時(shí)機(jī)進(jìn)行以保持一致性和效率。MOD 10**9 7 def dfs(pos, prev_color): ... total 0 for color in range(k): if color ! prev_color: total (total dfs(pos 1, color)) % MOD # 邊加邊模 dp[pos][prev_color] total return total同時(shí)要確保緩存中存儲的是取模后的值并且遞歸終點(diǎn)返回的1也要考慮取模雖然1 % MOD還是1。6.4 初始化與無效狀態(tài)處理對于使用數(shù)組緩存的情況初始化值如-1必須確保不會與任何有效結(jié)果混淆。如果有效結(jié)果可能為0或-1就需要選擇其他哨兵值或者使用一個(gè)單獨(dú)的visited布爾數(shù)組來記錄狀態(tài)是否已計(jì)算。另外要小心處理“無效狀態(tài)”。例如在“首尾不同色”問題中狀態(tài)(pos, prev_color, first_color)里的prev_color可能為-1起始時(shí)但first_color在起始時(shí)是未定義的。我們可以在遞歸函數(shù)開始時(shí)通過pos參數(shù)來區(qū)分是否需要檢查first_color或者用特殊的默認(rèn)值來表示“未定義”。7. 實(shí)戰(zhàn)演練解決一個(gè)綜合性的著色問題讓我們用一個(gè)稍微復(fù)雜點(diǎn)的例子來整合所有知識點(diǎn)。問題描述用k種顏色涂n個(gè)排成一列的格子。約束如下相鄰格子顏色不同。第一個(gè)格子和最后一個(gè)格子顏色也不能相同。顏色0最多只能使用limit次。我們將使用記憶化搜索來解決它。7.1 狀態(tài)定義這個(gè)問題結(jié)合了“首尾不同色”和“顏色次數(shù)限制”。我們需要跟蹤當(dāng)前處理到的位置pos(0到n)。前一個(gè)格子的顏色prev_color(-1到k-1)。第一個(gè)格子的顏色first_color(-1到k-1初始為-1表示未確定)。顏色0已經(jīng)使用的次數(shù)used_zero(0到limit)。因此狀態(tài)是一個(gè)四元組(pos, prev_color, first_color, used_zero)。7.2 狀態(tài)轉(zhuǎn)移與邊界處理遞歸終點(diǎn)(pos n): 檢查是否滿足“首尾不同色”約束。即如果first_color ! -1且prev_color first_color則此方案無效返回0否則返回1。當(dāng)前位置選擇顏色:遍歷所有顏色c(0 到 k-1)。約束1:c ! prev_color(除非prev_color -1即第一個(gè)格子)。約束3: 如果c 0則必須滿足used_zero 1 limit。狀態(tài)更新:new_used_zero used_zero (1 if c 0 else 0)new_first_color c if pos 0 else first_color(只有涂第一個(gè)格子時(shí)才確定first_color)遞歸調(diào)用dfs(pos1, c, new_first_color, new_used_zero)7.3 代碼實(shí)現(xiàn)與緩存由于狀態(tài)有四個(gè)維度且prev_color和first_color范圍是k1包含-1used_zero范圍是limit1使用四維數(shù)組可能代碼不夠清晰。這里我們使用functools.lru_cache配合元組作為鍵更為簡潔。from functools import lru_cache def solve_coloring(n, k, limit): MOD 10**9 7 lru_cache(maxsizeNone) def dfs(pos, prev_color, first_color, used_zero): # pos: 當(dāng)前要涂的格子索引 (0-based) # prev_color: 上一個(gè)格子的顏色-1表示沒有上一個(gè)起始狀態(tài) # first_color: 第一個(gè)格子的顏色-1表示尚未確定 # used_zero: 顏色0已經(jīng)使用的次數(shù) # 遞歸終點(diǎn)所有格子涂完 if pos n: # 檢查首尾顏色是否相同 if first_color ! -1 and prev_color first_color: return 0 # 違反約束2無效方案 return 1 # 找到一種合法方案 total 0 for color in range(k): # 約束1相鄰不能同色 (第一個(gè)格子跳過此檢查) if pos 0 and color prev_color: continue # 約束3顏色0使用次數(shù)限制 if color 0 and used_zero limit: continue # 計(jì)算新的狀態(tài) new_used_zero used_zero (1 if color 0 else 0) # 如果是第一個(gè)格子記錄其顏色 new_first_color color if pos 0 else first_color total (total dfs(pos 1, color, new_first_color, new_used_zero)) % MOD return total % MOD # 初始狀態(tài)從第0個(gè)格子開始前一個(gè)顏色無(-1)第一個(gè)顏色未定(-1)顏色0已使用0次 return dfs(0, -1, -1, 0) # 測試 n, k, limit 4, 3, 1 print(solve_coloring(n, k, limit)) # 輸出符合約束的方案數(shù)7.4 分析與優(yōu)化點(diǎn)這個(gè)解法直接、清晰但狀態(tài)空間是O(n * k * k * limit)。如果k和limit不大比如都10n在100左右是完全可行的。如果k很大我們可以注意到對于“顏色0”的特殊限制我們只額外跟蹤了它的使用次數(shù)而其他顏色是“無限制”且對稱的。這提示我們狀態(tài)中可以只區(qū)分“顏色0”和“非顏色0的其他顏色”而不是具體是哪種顏色從而將k的影響從狀態(tài)中部分剝離優(yōu)化狀態(tài)數(shù)量。這種基于對稱性的優(yōu)化是解決大規(guī)模計(jì)數(shù)問題的進(jìn)階技巧。通過這個(gè)綜合例子你應(yīng)該能感受到記憶化搜索就像一套“萬能模具”。面對新的約束我們主要的工作是設(shè)計(jì)出包含足夠信息的狀態(tài)表示然后遞歸關(guān)系往往可以比較直接地根據(jù)題意寫出來。剩下的就交給緩存去優(yōu)化效率。這種“定義狀態(tài)描述轉(zhuǎn)移緩存結(jié)果”的三段式思維是解決一大類組合計(jì)數(shù)問題的核心方法論。