记忆化搜索：从着色问题看复杂约束下的高效计数方案

1. 从一个“简单”的计数问题说起

最近在带新人刷算法题，遇到一个经典问题，它看起来人畜无害，却让不少初学者栽了跟头。题目大意是这样的：给你一个长度为n的格子，你需要用k种颜色去涂满它。但有一个限制：相邻的两个格子不能涂成相同的颜色。问一共有多少种不同的涂色方案？

乍一看，这不就是个排列组合题吗？第一个格子有k种选择，第二个格子不能和第一个相同，所以有k-1种选择，以此类推。总方案数不就是k * (k-1)^(n-1)吗？这个公式在n和k都不大的时候，确实能快速给出答案。很多新人做到这里就心满意足地提交了，然后……就收到了一个“Wrong Answer”。

问题出在哪？这个公式成立的前提是，颜色是“无限”的，或者说，我们每次选择时，可用的颜色数量只受“上一个格子颜色”这一个条件的约束。但在很多实际问题中，约束条件要复杂得多。比如，如果颜色数量k很小，或者题目增加了额外的限制，比如“首尾格子也不能同色”，甚至“某些特定位置的格子有固定颜色要求”，刚才那个简单的乘法原理就立刻失效了。这时，我们面对的就不再是一个有闭合公式的问题，而是一个需要系统化搜索所有可能状态的问题。

这就是“着色方案”类问题的核心：在满足一系列复杂约束条件下，计算所有可行的分配方案总数。当约束变得具体而微，暴力枚举所有可能性在数据规模稍大时就会变得不可能（时间复杂度是k^n的指数级）。此时，我们亟需一种更聪明的方法，而“记忆化搜索”正是为此而生的利器。它不是什么高深莫测的黑魔法，而是我们面对复杂状态空间时，一种化繁为简、避免重复劳动的朴素思想。接下来，我们就剥开这层外衣，看看它到底是怎么工作的，以及如何用它来优雅地解决那些看似棘手的计数问题。

2. 暴力搜索的困境与状态定义的艺术

在讨论记忆化搜索之前，我们必须先理解它所试图优化的对象——深度优先搜索（DFS）。对于着色问题，最直接的思路就是递归回溯：从第一个格子开始，尝试每一种可能的颜色，如果当前选择不违反约束（比如和左边格子颜色不同），就递归地去涂下一个格子。当所有格子都涂满时，就得到了一种合法方案，计数器加一。

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 # 初始调用，假设第一个格子左边没有格子，prev_color用-1表示 result = dfs(0, n, k, -1)

这段代码清晰易懂，但它有一个致命缺陷：存在大量重复计算。举个例子，假设n=5, k=3。当我们递归探索时，可能会先走颜色序列A->B->?这条路径，计算完后面所有的可能性。之后，在另一条分支里，我们可能又遇到了颜色序列C->B->?的状态。注意，此时虽然前两个格子的颜色不同（A和C），但第二个格子都是B，并且我们即将面对的是第三个格子。对于从“第三个格子开始，前一个颜色是B”这个子问题，它的答案是完全一样的，与第一个格子是A还是C无关！然而，我们的朴素DFS会傻乎乎地重新计算一遍。

这就是状态重叠。我们递归函数的本质，是在计算一个(当前位置, 前一个格子颜色)所确定的子问题的解。一旦这个二元组(pos, prev_color)确定了，无论通过哪条路径到达这个状态，后续的涂色方案数都是唯一确定的。如果我们能把这个结果存起来，下次再遇到相同的(pos, prev_color)时，直接返回结果，就能节省巨大的计算量。

所以，记忆化搜索的第一步，也是最重要的一步，就是精确定义“状态”。状态必须能唯一标识一个子问题，并且其数量是可控的。对于基本的相邻不同色问题，状态就是(pos, prev_color)。其中pos的范围是0到n（n表示已涂完，是递归终点），prev_color的范围是k种颜色，再加上一个表示“无前驱”的特殊值（比如-1）。因此，状态总数大约是(n+1) * (k+1)，这是一个多项式级别，远远小于指数级的k^n。

注意：状态定义并非一成不变。如果约束变成“首尾不能同色”，我们的状态就需要增加信息，比如变成(pos, prev_color, first_color)，因为最后一个格子的选择受第一个格子颜色的影响。定义状态的关键，在于找出哪些信息是决定后续选择所必需的、最小的信息集合。这需要根据具体问题的约束条件进行设计和提炼，是记忆化搜索中最具技巧性的部分。

3. 记忆化搜索的实现框架与细节打磨

理解了状态，实现记忆化搜索就水到渠成了。我们用一个缓存（通常是一个字典或数组）来存储已经计算过的状态结果。这个缓存结构的选择很有讲究。

1. 缓存数据结构的选择

• 字典（Dict/HashMap）：最通用和灵活。键（Key）是状态，值（Value）是结果。当状态比较复杂（比如包含多个离散变量）时，用字典很自然。例如，状态(pos, prev_color)可以转化为元组(pos, prev_color)作为键。
• 多维数组（List/Array）：当状态的所有维度都是整数且范围明确时，使用数组访问效率更高。例如，pos范围[0, n]，prev_color范围[-1, k-1]，我们可以建立一个(n+1) x (k+1)的二维数组dp，其中dp[pos][prev_color+1]存储结果（+1是为了将-1映射到索引0）。

在着色方案这类典型问题中，状态维度固定且范围小，使用数组是更优解。它不仅速度快，而且代码清晰。

2. 递归函数的改造我们将朴素的DFS函数改造成一个“有记忆”的DFS。

• 第一步：查缓存。在函数开始时，先检查当前状态是否已经计算过。如果是，直接返回缓存的结果。
• 第二步：递归计算。如果没计算过，则进行正常的递归逻辑，计算所有可能的选择并求和。
• 第三步：存缓存。在返回结果之前，将(当前状态, 计算结果)存入缓存。

以下是使用二维数组作为缓存的经典实现：

def count_colorings(n, k): # dp[pos][prev_color+1], 初始化所有值为-1表示未计算 # prev_color 从 -1 到 k-1，所以第二维大小是 k+1 dp = [[-1] * (k + 1) for _ in range(n + 1)] def dfs(pos, prev_color_idx): # prev_color_idx 是 prev_color 在dp数组中的索引 (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转换为“前一个颜色”的索引表示，用于比较 # 注意：prev_color_idx 是索引，真正的 prev_color = prev_color_idx - 1 actual_prev_color = prev_color_idx - 1 if color != actual_prev_color: # 递归，下一个位置的前一个颜色索引是 color + 1 total += dfs(pos + 1, color + 1) dp[pos][prev_color_idx] = total return total # 初始调用：从位置0开始，前一个颜色不存在，用索引0表示（即 actual_prev_color = -1） return dfs(0, 0) # 示例：5个格子，3种颜色，相邻不同色 print(count_colorings(5, 3)) # 输出应为 3 * 2^4 = 48，可以用公式验证

3. 边界条件与初始化递归的终点（pos == n）通常返回 1（表示找到一种完整方案）。缓存数组的初始化值必须是一个不会出现在正常结果中的值（如-1），用以区分“未计算”和“计算结果为0”（后者在某些问题中是合法结果，表示无解）。

4. 复杂度分析

• 时间复杂度：由于每个状态(pos, prev_color)最多只计算一次，每次计算需要遍历k种颜色，所以总时间复杂度为O(n * k * k)？等等，仔细看内层循环。对于每个状态，我们循环k次，每次递归调用是 O(1) 的查表或计算。因此，准确的时间复杂度是O(状态数 * 每个状态的计算成本) = O(n * k * 1) = O(n * k)。这里的k是颜色数，通常是个常数或者不大的数，因此算法是线性或近似线性的，效率极高。
• 空间复杂度：主要是缓存数组dp的开销，为O(n * k)，以及递归调用栈的深度O(n)。

实操心得：在实现时，我强烈建议将“状态”到“缓存索引”的映射关系单独写成一个清晰的函数或注释。比如get_index(prev_color)。这能极大减少因为下标转换错误导致的Bug，尤其是在状态变量有特殊值（如-1）的时候。另外，对于结果可能非常大的计数问题（比如方案数可能超过64位整数范围），要在题目要求下及时取模，并且在存入缓存和返回结果前都要取模，保证一致性。

4. 从经典到变种：应对更复杂的约束条件

记忆化搜索的强大之处在于其灵活性。当问题的约束条件发生变化时，我们通常不需要推翻重来，而只需调整“状态定义”和“状态转移”逻辑。下面我们通过几个变种问题来体会这一点。

4.1 变种一：首尾格子也不能同色这是“相邻不同色”问题的经典加强版。此时，最后一个格子（第n-1个）的颜色不仅不能和它左边的格子（第n-2个）相同，还不能和第一个格子相同。

状态定义的升级：原来的状态(pos, prev_color)不足以决定最后一个格子的选择，因为它缺少了“第一个格子颜色”的信息。因此，我们需要将“第一个格子的颜色”也纳入状态。定义状态为(pos, prev_color, first_color)。其中first_color在递归开始时就确定下来，并一路传递下去。

状态转移的调整：在递归涂色时：

• 当pos == 0（涂第一个格子）：遍历所有k种颜色作为first_color，同时这个颜色也是prev_color。
• 当pos == n-1（涂最后一个格子）：遍历颜色时，除了要满足color != prev_color，还必须满足color != first_color。
• 其他位置：和之前一样，只需满足color != prev_color。

缓存维度：状态变成了三维(pos, prev_color, first_color)，缓存数组的大小变为(n) * (k) * (k)。虽然空间变大了，但相对于指数爆炸，这依然是完全可以接受的。

4.2 变种二：颜色使用次数限制假设每种颜色最多只能使用m次。这在实际场景中很常见，比如有限的颜料库存。

状态定义的升级：此时，仅仅知道前一个颜色是什么不够了，我们还需要知道每种颜色还剩多少使用次数。一种直观的状态定义是(pos, prev_color, color_used_tuple)，其中color_used_tuple是一个长度为k的元组，记录每种颜色已使用的次数。但这样状态空间会非常大（n * k * (m+1)^k）。

优化思路：对于计数问题，我们往往不需要知道每种颜色具体用了多少次，而只需要知道“剩余使用次数”的模式。如果所有颜色的限制次数m相同，那么问题可以简化为：在涂到某个位置时，有多少种颜色已经用满了m次，有多少种颜色用了m-1次……但这依然复杂。一个更实用的方法是，当k和m不大时，可以使用状态压缩。用一个k位的整数（比特位）来表示哪些颜色已经用尽了次数，或者用一个整数数组来记录使用次数，并将整个数组作为字典的键（虽然效率会降低）。这体现了记忆化搜索的另一个维度：当状态本身复杂时，我们可以利用哈希表（字典）的灵活性来存储。

4.3 变种三：格子分组着色（图着色问题的简化）问题升级为：格子之间不是简单的线性关系，而是一个一般的图。每个节点（格子）需要着色，有边相连的节点不能同色。这就是经典的图着色问题，是NP难的。但对于特定的、树状或稀疏的图，记忆化搜索结合树形DP仍然可以高效解决。

状态定义：对于树形结构，我们通常在树上进行DFS。状态可以定义为(node, parent_color)，表示在以node为根的子树中，当node的父节点颜色为parent_color时，该子树的着色方案数。然后通过递归合并子节点的结果来计算当前节点的方案数。

踩坑实录：在处理复杂约束时，最容易犯的错误是状态定义遗漏了关键信息。我曾在一个比赛中遇到一个问题，要求“任意两个距离为2的格子也不能同色”。我最初只定义了(pos, prev_color)，结果总是少算。后来才意识到，距离为2意味着当前格子不能和它前面第2个格子同色。因此，状态必须包含前两个格子的颜色信息，即(pos, color_of_pos_minus_1, color_of_pos_minus_2)。这个教训让我明白，定义状态时要像侦探一样，问自己：“要唯一确定从现在开始的所有未来可能性，最少需要知道过去的哪些信息？”

5. 记忆化搜索 vs. 动态规划：思维路径的异同

很多人会把记忆化搜索和动态规划（DP）等同起来，称其为“递归形式的DP”。这种说法有一定道理，但两者在思维起点和实现方式上有着微妙的区别，理解这些区别能帮助你更好地选择工具。

5.1 思维路径的对比

• 记忆化搜索（Memoization）：思维是自顶向下的。你从要解决的原问题（如f(0, -1)）开始思考：“要解决我的问题，我需要先解决哪些子问题？”然后递归地去解决这些子问题，并用缓存避免重复。它的思路更符合人类面对复杂问题的自然分解过程——分而治之。
• 动态规划（Dynamic Programming）：思维是自底向上的。你需要先确定所有子问题的计算顺序（通常是较小的、基础的状态先计算），然后通过迭代循环，从小问题逐步推导出大问题的解。这需要更强的“全局”状态转移视角。

对于着色方案问题，记忆化搜索的思维是：“我想知道从第0个格子开始涂有几种方案。那我先试试涂第一种颜色，然后问题就变成了‘从第1个格子开始，且前一个颜色是第一种颜色’有几种方案。我去计算这个子问题……”而动态规划则会先计算“最后一个格子怎么涂”，然后倒推回来，或者从第一个格子开始正推。

5.2 实现形式的对比我们以基础着色问题为例，看看两者的代码实现。

记忆化搜索（递归）：如上文所示，代码直观反映了递归关系。

动态规划（迭代）：我们需要定义dp[i][c]表示“涂完前i个格子，并且第i个格子（最后一个）颜色是c的方案总数”。

• 状态转移：dp[i][c] = sum(dp[i-1][c'])，其中c'是所有不等于c的颜色。因为第i个格子涂c，那么第i-1个格子可以是任何非c的颜色。
• 初始化：dp[0][c] = 1（对于第一个格子，每种颜色都是一种方案）。
• 最终答案：sum(dp[n-1][c])对所有颜色c求和。

def count_colorings_dp(n, k): if n == 0: return 0 # dp[i][c]: 前i个格子已涂完，且第i个格子颜色为c的方案数 (i从0开始) dp = [[0] * k for _ in range(n)] # 初始化第一个格子 for c in range(k): dp[0][c] = 1 # 递推 for i in range(1, n): for c in range(k): # 当前格子涂c，上一个格子可以涂任何非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 total

5.3 如何选择？

• 优先考虑记忆化搜索：当状态转移关系不那么直观，或者存在复杂的依赖关系（比如在树上），记忆化搜索更容易思考和实现。你只需要写出递归关系，让缓存去处理重复子问题。
• 考虑动态规划：当问题有明显的线性顺序，且状态转移方程清晰简单时，DP的迭代形式通常效率稍高（避免了递归调用开销），并且不容易出现栈溢出（对于深度很大的递归）。
• 一个实用的建议：先尝试用记忆化搜索的思路去思考和解决问题。写出递归函数。如果发现性能或栈深度有问题，再考虑是否能转化为等价的、自底向上的动态规划。很多时候，记忆化搜索是探索DP状态转移方程的绝佳跳板。

个人经验：在竞赛或面试中，如果时间紧迫，我通常会首选记忆化搜索。因为它更不容易出错，思维负担小。只要确保状态定义正确、缓存生效，基本就能拿到分数。而自底向上的DP，一旦递推顺序或初始化写错，调试起来可能更费时间。当然，对于状态空间巨大、需要滚动数组优化空间的情况，就必须使用迭代DP了。

6. 性能优化与边界陷阱

即使使用了记忆化搜索，如果不注意细节，依然可能掉入性能或正确性的陷阱。这里分享几个关键点。

6.1 缓存键的设计与哈希效率当使用字典（Python的dict或functools.lru_cache）时，状态的哈希效率至关重要。最常用的方法是将状态转换为元组（Tuple）。但要注意：

• 如果状态中包含列表（List），必须先转换为元组，因为列表是不可哈希的。
• 对于整数状态，直接使用元组即可。对于复杂对象，可以考虑使用字符串编码或自定义哈希函数。

在Python中，使用@lru_cache(maxsize=None)装饰器可以极简地实现记忆化，它自动将函数参数作为缓存键。这对于原型设计和快速验证非常方便。

from functools import lru_cache @lru_cache(maxsize=None) 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 total

6.2 递归深度限制Python默认的递归深度限制（通常为1000）对于n较大的问题可能不够。对于线性递归深度为n的问题，当n超过1000时，需要手动设置递归深度或改用迭代DP。

import sys sys.setrecursionlimit(10000) # 设置为一个更大的值

但更根本的解决方法是，评估问题是否必须深度递归。像着色方案这种问题，递归深度等于格子数n，如果n达到10^5级别，即使解除限制，递归调用栈的开销也极大，且有栈溢出风险。此时，必须使用迭代的动态规划。

6.3 大数取模的处理方案数往往非常巨大，题目通常要求对某个大数MOD（如10^9+7）取模。这里有一个极易出错的细节：必须在每一次加法运算后立即取模，而不是最后才取模。因为中间结果可能已经溢出（即使在Python这种大整数语言中，取模操作本身也应在合理时机进行以保持一致性和效率）。

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

同时，要确保缓存中存储的是取模后的值，并且递归终点返回的1也要考虑取模（虽然1 % MOD还是1）。

6.4 初始化与无效状态处理对于使用数组缓存的情况，初始化值（如-1）必须确保不会与任何有效结果混淆。如果有效结果可能为0或-1，就需要选择其他哨兵值，或者使用一个单独的visited布尔数组来记录状态是否已计算。

另外，要小心处理“无效状态”。例如，在“首尾不同色”问题中，状态(pos, prev_color, first_color)里的prev_color可能为-1（起始时），但first_color在起始时是未定义的。我们可以在递归函数开始时，通过pos参数来区分是否需要检查first_color，或者用特殊的默认值来表示“未定义”。

7. 实战演练：解决一个综合性的着色问题

让我们用一个稍微复杂点的例子来整合所有知识点。问题描述：用k种颜色涂n个排成一列的格子。约束如下：

1. 相邻格子颜色不同。
2. 第一个格子和最后一个格子颜色也不能相同。
3. 颜色0最多只能使用limit次。

我们将使用记忆化搜索来解决它。

7.1 状态定义这个问题结合了“首尾不同色”和“颜色次数限制”。我们需要跟踪：

1. 当前处理到的位置pos(0到n)。
2. 前一个格子的颜色prev_color(-1到k-1)。
3. 第一个格子的颜色first_color(-1到k-1，初始为-1表示未确定)。
4. 颜色0已经使用的次数used_zero(0到limit)。

因此，状态是一个四元组：(pos, prev_color, first_color, used_zero)。

7.2 状态转移与边界处理

• 递归终点(pos == n): 检查是否满足“首尾不同色”约束。即，如果first_color != -1且prev_color == first_color，则此方案无效，返回0；否则返回1。
• 当前位置选择颜色:
  • 遍历所有颜色c(0 到 k-1)。
  • 约束1:c != prev_color(除非prev_color == -1，即第一个格子)。
  • 约束3: 如果c == 0，则必须满足used_zero + 1 <= limit。
• 状态更新:
  • new_used_zero = used_zero + (1 if c == 0 else 0)
  • new_first_color = c if pos == 0 else first_color(只有涂第一个格子时才确定first_color)
  • 递归调用：dfs(pos+1, c, new_first_color, new_used_zero)

7.3 代码实现与缓存由于状态有四个维度，且prev_color和first_color范围是k+1（包含-1），used_zero范围是limit+1，使用四维数组可能代码不够清晰。这里我们使用functools.lru_cache配合元组作为键，更为简洁。

from functools import lru_cache def solve_coloring(n, k, limit): MOD = 10**9 + 7 @lru_cache(maxsize=None) def dfs(pos, prev_color, first_color, used_zero): # pos: 当前要涂的格子索引 (0-based) # prev_color: 上一个格子的颜色，-1表示没有上一个（起始状态） # first_color: 第一个格子的颜色，-1表示尚未确定 # used_zero: 颜色0已经使用的次数 # 递归终点：所有格子涂完 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：相邻不能同色 (第一个格子跳过此检查) if pos > 0 and color == prev_color: continue # 约束3：颜色0使用次数限制 if color == 0 and used_zero >= limit: continue # 计算新的状态 new_used_zero = used_zero + (1 if color == 0 else 0) # 如果是第一个格子，记录其颜色 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 # 初始状态：从第0个格子开始，前一个颜色无(-1)，第一个颜色未定(-1)，颜色0已使用0次 return dfs(0, -1, -1, 0) # 测试 n, k, limit = 4, 3, 1 print(solve_coloring(n, k, limit)) # 输出符合约束的方案数

7.4 分析与优化点这个解法直接、清晰，但状态空间是O(n * k * k * limit)。如果k和limit不大（比如都<=10），n在100左右，是完全可行的。如果k很大，我们可以注意到，对于“颜色0”的特殊限制，我们只额外跟踪了它的使用次数，而其他颜色是“无限制”且对称的。这提示我们，状态中可以只区分“颜色0”和“非颜色0的其他颜色”，而不是具体是哪种颜色，从而将k的影响从状态中部分剥离，优化状态数量。这种基于对称性的优化是解决大规模计数问题的进阶技巧。

通过这个综合例子，你应该能感受到，记忆化搜索就像一套“万能模具”。面对新的约束，我们主要的工作是设计出包含足够信息的状态表示，然后递归关系往往可以比较直接地根据题意写出来。剩下的，就交给缓存去优化效率。这种“定义状态，描述转移，缓存结果”的三段式思维，是解决一大类组合计数问题的核心方法论。