ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

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

记忆化搜索:从着色问题看复杂约束下的高效计数方案 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)这段代码清晰易懂但它有一个致命缺陷存在大量重复计算。举个例子假设n5, k3。当我们递归探索时可能会先走颜色序列A-B-?这条路径计算完后面所有的可能性。之后在另一条分支里我们可能又遇到了颜色序列C-B-?的状态。注意此时虽然前两个格子的颜色不同A和C但第二个格子都是B并且我们即将面对的是第三个格子。对于从“第三个格子开始前一个颜色是B”这个子问题它的答案是完全一样的与第一个格子是A还是C无关然而我们的朴素DFS会傻乎乎地重新计算一遍。这就是状态重叠。我们递归函数的本质是在计算一个(当前位置, 前一个格子颜色)所确定的子问题的解。一旦这个二元组(pos, prev_color)确定了无论通过哪条路径到达这个状态后续的涂色方案数都是唯一确定的。如果我们能把这个结果存起来下次再遇到相同的(pos, prev_color)时直接返回结果就能节省巨大的计算量。所以记忆化搜索的第一步也是最重要的一步就是精确定义“状态”。状态必须能唯一标识一个子问题并且其数量是可控的。对于基本的相邻不同色问题状态就是(pos, prev_color)。其中pos的范围是0到nn表示已涂完是递归终点prev_color的范围是k种颜色再加上一个表示“无前驱”的特殊值比如-1。因此状态总数大约是(n1) * (k1)这是一个多项式级别远远小于指数级的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]我们可以建立一个(n1) x (k1)的二维数组dp其中dp[pos][prev_color1]存储结果1是为了将-1映射到索引0。在着色方案这类典型问题中状态维度固定且范围小使用数组是更优解。它不仅速度快而且代码清晰。2. 递归函数的改造我们将朴素的DFS函数改造成一个“有记忆”的DFS。第一步查缓存。在函数开始时先检查当前状态是否已经计算过。如果是直接返回缓存的结果。第二步递归计算。如果没计算过则进行正常的递归逻辑计算所有可能的选择并求和。第三步存缓存。在返回结果之前将(当前状态, 计算结果)存入缓存。以下是使用二维数组作为缓存的经典实现def count_colorings(n, k): # dp[pos][prev_color1], 初始化所有值为-1表示未计算 # 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数组中的索引 (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 * (m1)^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 total5.3 如何选择优先考虑记忆化搜索当状态转移关系不那么直观或者存在复杂的依赖关系比如在树上记忆化搜索更容易思考和实现。你只需要写出递归关系让缓存去处理重复子问题。考虑动态规划当问题有明显的线性顺序且状态转移方程清晰简单时DP的迭代形式通常效率稍高避免了递归调用开销并且不容易出现栈溢出对于深度很大的递归。一个实用的建议先尝试用记忆化搜索的思路去思考和解决问题。写出递归函数。如果发现性能或栈深度有问题再考虑是否能转化为等价的、自底向上的动态规划。很多时候记忆化搜索是探索DP状态转移方程的绝佳跳板。个人经验在竞赛或面试中如果时间紧迫我通常会首选记忆化搜索。因为它更不容易出错思维负担小。只要确保状态定义正确、缓存生效基本就能拿到分数。而自底向上的DP一旦递推顺序或初始化写错调试起来可能更费时间。当然对于状态空间巨大、需要滚动数组优化空间的情况就必须使用迭代DP了。6. 性能优化与边界陷阱即使使用了记忆化搜索如果不注意细节依然可能掉入性能或正确性的陷阱。这里分享几个关键点。6.1 缓存键的设计与哈希效率当使用字典Python的dict或functools.lru_cache时状态的哈希效率至关重要。最常用的方法是将状态转换为元组Tuple。但要注意如果状态中包含列表List必须先转换为元组因为列表是不可哈希的。对于整数状态直接使用元组即可。对于复杂对象可以考虑使用字符串编码或自定义哈希函数。在Python中使用lru_cache(maxsizeNone)装饰器可以极简地实现记忆化它自动将函数参数作为缓存键。这对于原型设计和快速验证非常方便。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默认的递归深度限制通常为1000对于n较大的问题可能不够。对于线性递归深度为n的问题当n超过1000时需要手动设置递归深度或改用迭代DP。import sys sys.setrecursionlimit(10000) # 设置为一个更大的值但更根本的解决方法是评估问题是否必须深度递归。像着色方案这种问题递归深度等于格子数n如果n达到10^5级别即使解除限制递归调用栈的开销也极大且有栈溢出风险。此时必须使用迭代的动态规划。6.3 大数取模的处理方案数往往非常巨大题目通常要求对某个大数MOD如10^97取模。这里有一个极易出错的细节必须在每一次加法运算后立即取模而不是最后才取模。因为中间结果可能已经溢出即使在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个排成一列的格子。约束如下相邻格子颜色不同。第一个格子和最后一个格子颜色也不能相同。颜色0最多只能使用limit次。我们将使用记忆化搜索来解决它。7.1 状态定义这个问题结合了“首尾不同色”和“颜色次数限制”。我们需要跟踪当前处理到的位置pos(0到n)。前一个格子的颜色prev_color(-1到k-1)。第一个格子的颜色first_color(-1到k-1初始为-1表示未确定)。颜色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(pos1, c, new_first_color, new_used_zero)7.3 代码实现与缓存由于状态有四个维度且prev_color和first_color范围是k1包含-1used_zero范围是limit1使用四维数组可能代码不够清晰。这里我们使用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: 当前要涂的格子索引 (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不大比如都10n在100左右是完全可行的。如果k很大我们可以注意到对于“颜色0”的特殊限制我们只额外跟踪了它的使用次数而其他颜色是“无限制”且对称的。这提示我们状态中可以只区分“颜色0”和“非颜色0的其他颜色”而不是具体是哪种颜色从而将k的影响从状态中部分剥离优化状态数量。这种基于对称性的优化是解决大规模计数问题的进阶技巧。通过这个综合例子你应该能感受到记忆化搜索就像一套“万能模具”。面对新的约束我们主要的工作是设计出包含足够信息的状态表示然后递归关系往往可以比较直接地根据题意写出来。剩下的就交给缓存去优化效率。这种“定义状态描述转移缓存结果”的三段式思维是解决一大类组合计数问题的核心方法论。
返回列表