P1036选数：DFS组合枚举与素数判定的经典入门题

第一次看到 P1036 这道题的时候我刚学递归没多久脑子里只有一个念头“从 n 个数里选 k 个”这种问题只能靠全排列硬莽。结果被 WA 教育了好几次才发现组合枚举里藏着不少细节素数判定也不是写个 for 循环就万事大吉。这道题是 NOIP 2002 普及组的第二题题目叫“选数”二十多年过去了它依然是很多 OIer 入门 DFS 组合枚举和素数判定的第一道坎。今天把这题从头到尾拆一遍说说 DFS 枚举组合的原理、试除法的边界处理以及那些真正会让人卡住的小坑。1. 先把题意看清选数到底在选什么1.1 输入输出与样例题目简洁得不像话给你 n 个整数以及一个整数 kk n要求从这 n 个数中任选 k 个相加得到一系列和统计其中有多少个和是素数输出这个数量。输入格式就两行4 3 3 7 12 19输出一个整数。样例的输出是 1因为 4 个数里选 3 个组合一共只有 C(4, 3) 4 种3 7 12 22不是素数3 7 19 29是素数3 12 19 34不是素数7 12 19 38不是素数所以答案就是 1。看起来很简单对吧真正写起来问题就来了怎么保证枚举的是“组合”而不是“排列”怎么知道 29 是不是素数这些才是这道题真正的考点。1.2 数据范围决定了算法走向我看过很多同学拿到题第一件事就是琢磨“是不是得用 DP 或者什么高端优化”。说实话在竞赛里第一步永远是看数据范围数据范围直接决定了你能用什么复杂度的算法。原题的数据范围我记得很清楚1 ≤ n ≤ 201 ≤ k n1 ≤ xi ≤ 5000000n 只有 20这是什么概念组合数 C(20, 10) 是 184756也就是说就算我把所有情况全部枚举一遍最多也就 18 万多种组合。这个量级对于现代评测机来说连“压力”都算不上。所以这道题的本质就是一个暴力枚举题问题只在于你怎么把这个“暴力”写得优雅、不重复、不漏解。数据范围里的 xi 上限 5000000 同样有讲究。k 最坏接近 19那么最大的和约为 19 × 5000000 95000000接近 1 亿。判断一个接近一亿的数是不是素数用试除法只需要枚举到 sqrt(95000000) ≈ 9747 次这个代价在 18 万种组合下完全可控。换句话说题目设计者就是故意给了这个范围让你用 DFS 试除法就能舒服地 AC。2. 算法设计DFS 组合枚举为什么是正解2.1 组合与排列关键的思维转换我第一次写这题时用的是标准全排列模板void dfs(int step) { if (step k) { // 判断求和是否是素数 return; } for (int i 0; i n; i) { if (!used[i]) { used[i] true; // 记录选择 dfs(step 1); used[i] false; } } }这份代码跑出来的结果对于样例是对的但自己构造几个数据就发现答案偏大。原因很简单全排列把“选择 3、7、19”和“选择 19、7、3”当成了两种完全不同的方案而题目显然只算一种。全排列枚举出的方案数会是组合数的 k! 倍也就是排列考虑了顺序组合不管顺序。“组合”的本质是从 n 个元素里选出一个大小为 k 的无序集合。无序意味着我们只需要关心“选了哪些数”不需要关心“先选谁后选谁”。所以一个很自然的想法就诞生了在递归枚举时强制要求后选的数必须排在先选的数后面这样天然就不会出现顺序不同的重复组合。2.2 start 参数让枚举天然不重复DFS 组合枚举的经典写法关键在于给递归函数加一个“起始下标”参数。核心代码可以短到只有十几行void dfs(int step, int start, int sum) { if (step k) { if (isPrime(sum)) ans; return; } for (int i start; i n; i) { dfs(step 1, i 1, sum a[i]); } }这里的三个参数各司其职step当前已经选了几个数对应递归深度。start下一层递归可以从哪个下标开始选这是组合不重复的关键。sum当前已选数字的累加和直接通过参数传递不需要全局变量。当递归到 step k 时说明已经选满 k 个数此时 sum 就是这 k 个数的和调用素数判定即可。为什么start能保证不重复因为每次递归调用传的是i 1也就是说下一层只能从当前选中位置的下一个位置开始挑。比如选择了下标 2 的数之后后面只能选下标 3、4、5……这样就保证了任意一种组合只会被枚举到一次选中的下标序列永远是严格递增的。严格递增就等价于“无序”因为任何一组无序选法都可以按下标从小到大排序而且排序方式唯一。2.3 递归全过程模拟空讲不太好懂我用 n 4, k 3, a {3, 7, 12, 19} 手动把递归树走一遍你就能直观感受到 start 参数是怎么工作的。第一层step 0start 0sum 0循环 i 从 0 到 3选 a[0] 3进入 dfs(1, 1, 3)i 1选 7进入 dfs(2, 2, 10)i 2选 12进入 dfs(3, 3, 22)22 不是素数i 3选 19进入 dfs(3, 4, 29)29 是素数ans 1i 2选 12进入 dfs(2, 3, 15)i 3选 19进入 dfs(3, 4, 34)34 不是素数i 3选 19进入 dfs(2, 4, 22)此时 start 4循环直接结束选不满 3 个自然返回选 a[1] 7进入 dfs(1, 2, 7)i 2选 12进入 dfs(2, 3, 19)i 3选 19进入 dfs(3, 4, 38)38 不是素数i 3选 19进入 dfs(2, 4, 26)选不满结束选 a[2] 12进入 dfs(1, 3, 12)i 3选 19进入 dfs(2, 4, 31)选不满结束选 a[3] 19进入 dfs(1, 4, 19)直接结束注意那些“选不满 k 个就结束”的情况。比如第一层的 i 3 时选了最后一个下标剩余可选位置已经不足但代码不会报错——for 循环条件i n自然不成立递归就会返回。这也是组合枚举和全排列一个很大的区别全排列用了 used 数组递归深度必须达到 k组合枚举用 start 限制搜索空间递归可能提前“无路可走”但这种情况完全不影响正确性。2.4 顺带聊聊其他枚举组合的姿势除了 DFS start 参数组合枚举还有其他几种常见写法各有适用场景二进制枚举子集用一个 mask 从 0 枚举到 2^n - 1统计二进制位中 1 的个数等于 k 的 mask然后累加对应位置的数。这种方法写起来短适合 n ≤ 20 的场合。缺点是每次都要重新遍历 n 个数计算和而且无法在枚举过程中做剪枝。next_permutation 法构造一个数组前 k 个是 1、后面是 0用 next_permutation 生成所有排列然后按 1 的位置求和。代码相对绕而且 C 的 next_permutation 对于有重复元素的数组会跳过重复排列有时候反而不直观。递归“选/不选”分叉每个数有两种决策选或不选。需要额外记录选了几个数走到末尾时判断 cnt 是否等于 k。这种写法本质上就是子集枚举思路很清晰但搜索树深度是 n分支是 2对于本题也能过。我自己最推荐的就是 start 参数版 DFS因为它最接近“组合”的数学定义搜索规模正好是 C(n, k)没有浪费任何状态。这个模板理解透之后后面做“8 皇后”、“数独”、“全排列”都能举一反三。3. 素数判定试除法里面的细节3.1 为什么这道题不需要筛素数表很多刚接触算法的同学看到“素数”两个字第一反应就是打表先写个埃氏筛把所有素数筛出来。这当然是一种思路但我们要先想想有没有必要。这道题的最大和大约是 95000000接近 1 亿。如果你开一个 1 亿大小的 bool 数组来标记素数内存大约是 100MB不少评测机会出问题。而且埃氏筛需要 O(n log log n) 的预处理时间对于只要判断 18 万次素数的题目来说属于杀鸡用了牛刀。更关键的是DFS 过程中产生的 sum 不是连续的也不是递增的你没法像离线查询那样方便地复用筛法结果。至于 Miller-Rabin 这种工业级素性测试更是完全没必要。对于一个最大不超过 1 亿的数试除法就是最简单、最不容易出错的选择。3.2 一份不容易写错的试除函数标准的试除法大家都知道从 2 开始枚举看 x 能否被整除。但边界条件的处理是很多人栽跟头的地方。bool isPrime(int x) { if (x 2) return false; for (int i 2; i * i x; i) { if (x % i 0) return false; } return true; }几个关键点if (x 2) return false;是第一道防线。1 不是素数0 和负数也不是。题目里的 xi 都是正整数sum 一定大于 0但谁也不能保证你以后不会拿这个函数去判断别的数写上这个判断没有坏处。循环条件是i * i x等价于i sqrt(x)。当你写i * i x的时候i 的平方不会超过 x 的平方根由于 x 最大不到 1 亿i 最大也就 10000 左右i * i撑死 1 亿远远不会超过 int 的范围所以不存在溢出风险。为什么不是i sqrt(x)因为sqrt()返回的是浮点数浮点比较在某些极端边界值下可能因为精度问题出错。比如 sqrt(9) 理论上是 3.0但某些编译器实现可能得到 2.9999999999999996导致循环少判断一位。虽然这种概率很低但竞赛里一行代码的差值可能决定一次提交的成败所以用整数乘法判断最稳妥。3.3 小优化把循环次数减半试除法可以在常数级别上再优化一把。因为除 2 以外的所有偶数都不可能是素数能被 2 整除所以我们可以先单独判断偶数的情况然后再从 3 开始步长为 2 地枚举奇数因子。bool isPrime(int x) { if (x 2) return false; if (x % 2 0) return x 2; for (int i 3; i * i x; i 2) { if (x % i 0) return false; } return true; }这个版本有一个小细节if (x % 2 0) return x 2;。当 x 2 时2 % 2 0同时 x 2所以返回 true当 x 4、6、8 这些偶数时x % 2 0 且 x ! 2返回 false。这个写法看着有点绕但实际非常简洁用一句话同时处理了“偶数且为 2”和“偶数且非 2”两种情况。为什么可以放心跳过所有偶数因子如果一个数 x 能被偶数 d 整除那么 x 一定也能被 2 整除。既然 x 是奇数已经排除了偶数情况那么它就不可能存在偶数因子。所以枚举奇数因子就够了。这个优化能把试除法的耗时减半在本题环境下属于锦上添花但在时间复杂度卡得比较死的题里这种小习惯能救你一命。4. 完整 AC 代码与复杂度分析4.1 参考实现带注释把前面所有思路组合到一起就是一份可以 AC 的完整代码。我用的纯 C没有任何花哨的库你拿到什么环境下都能编译运行#include bits/stdc.h using namespace std; const int MAXN 25; int n, k; int a[MAXN]; int ans; bool isPrime(int x) { if (x 2) return false; if (x % 2 0) return x 2; for (int i 3; i * i x; i 2) { if (x % i 0) return false; } return true; } void dfs(int step, int start, int sum) { if (step k) { if (isPrime(sum)) ans; return; } for (int i start; i n; i) { dfs(step 1, i 1, sum a[i]); } } int main() { ios::sync_with_stdio(false); cin.tie(0); cin n k; for (int i 0; i n; i) { cin a[i]; } dfs(0, 0, 0); cout ans \n; return 0; }4.2 递归边界和参数设计这份代码里最核心的就是 dfs 函数的三参数设计。很多人会问为什么不把 sum 设成全局变量其实完全可以但设成全局变量后就必须在递归返回时做“回溯还原”操作sum a[i]; dfs(step 1, i 1); sum - a[i];如果不写最后那行sum - a[i]那么递归返回后当前层的 sum 已经被“污染”了下一次循环就会在一个错误的基础值上累加。而把 sum 作为值传递的参数每个递归分支都有独立的 sum 副本天然不会相互干扰也就省去了回溯步骤。这也是我个人强烈推荐值传递的原因少一行回溯就少一个出错的机会。类似的道理也适用于 step 参数。有的同学把 step 也设成全局变量那么递归前后就要 step 和 step--代码看起来就没那么清爽。竞赛里写代码怎么不容易错怎么来值传递就是最不容易错的方式。4.3 时间复杂度到底够不够我们来算一笔账。最坏情况下n 20 时组合数最大出现在 k 取 10 左右C(20, 10) 184756。每次选满 k 个数后要做一次素数判定每次 isPrime 的最坏代价大约是 sqrt(95000000) ≈ 9747 次取模运算。两者相乘理论上限大约是 18 亿次运算。这个数字听起来很吓人但你仔细想想实际情况第一只有选满 k 个才会调用 isPrime而搜索树的内部节点远多于叶子节点内部节点不做素数判定第二不是每个 sum 都接近最大值大多数和远小于上限试除次数随之减少第三我们用了偶数优化实际循环次数再减半第四C(20, 10) 是理论最大k 很可能不是 10。综合下来实际评测时运行时间通常在 10 毫秒级别就算评测机再慢这道题 1 秒的时限也绰绰有余。空间复杂度就更不用担心了。递归深度最多 k 层k n ≤ 20数组 a 长度 20整个程序占用的内存可以忽略不计O(n k) 级别。5. 踩坑记录我在这道题上交过的学费5.1 把组合写成了全排列这个坑我在 2.1 节里提过。具体表现就是样例能过但自己构造的数据答案偏大。比如 n 4, k 3数据还是 {3, 7, 12, 19}如果用全排列你会把 3719、7319、1937 当成三个不同方案最后答案翻 6 倍3! 6。如果你发现你的答案正好是正确答案的 k! 倍那基本可以断定问题出在组合/排列没分清。5.2 素数判定漏掉边界漏掉x 2的判断isPrime(1) 会返回 true因为 for 循环从 2 开始1 % 2 1循环条件i * i x直接不成立函数就返回 true 了。如果某组数据的和恰好是 1题目中 xi 是正整数所以 k 个数相加最小也是 k理论上不会出现 1但你调用其他测试数据时就可能炸就会导致误判。5.3 用 sqrt 写循环导致变慢有些版本会写成for (int i 2; i sqrt(x); i)。这写法的问题有两个第一每次循环都要调用 sqrt() 浮点库函数速度比整数乘法慢一个量级第二浮点数精度在极端情况下可能出问题。如果你习惯了这个写法建议改掉以后遇到时间复杂度敏感的题这个习惯可能让你 TLE。正确写法是for (int i 2; i * i x; i)。如果担心 i * i 溢出可以写成i x / i这也是常见安全写法。5.4 全局 sum 没回溯这个坑很隐蔽尤其是刚开始学 DFS 的时候。把 sum 设成全局变量进入递归时累加返回时忘记还原会导致同一个值被反复使用最后结果是各种组合的混乱拼接。调试方法也很简单在 dfs 函数入口输出 step、start、sum 三个变量肉眼检查累加值是否符合预期。5.5 读入顺序和调试技巧还有同学会把cin n k手滑写成cin k n。由于题目保证 k n而数据又是 k 在左边这种错误在简单样例里不一定暴露因为 k 和 n 数值都比较小。建议养成习惯写输入后先用样例跑一遍再把样例的 n、k 对调构造一个自测数据能有效排查这种低级错误。如果实在找不到 bug我的土办法是暴力对拍。写一个三重循环或者用 next_permutation 枚举所有组合用同样的 isPrime 函数求和判断和 DFS 版本跑随机数据比对输出。随机生成 100 组小数据两边输出不一致就一定有问题逐步缩小区间就能定位到具体哪一步出错。6. 做完这题之后还能往哪走6.1 老题对今天 CSP-J 的意义NOIP 普及组现在已经改名为 CSP-J也就是面向初中生的非专业级软件能力认证。很多同学觉得 2002 年的题目已经过时了但你去翻 CSP-J 的真题DFS、组合枚举、素数判定依然反复出现。CSP-J 2023 年的完善程序题里就有递归搜索的影子CSP-J 2022 的阅读程序也考了类似的逻辑。这套知识体系二十年没有变过变的只是题目的外壳。把 P1036 这种基础题吃透等于给 CSP-J 的搜索题和大模拟题打下一个扎实的地基。6.2 换个思路重写这题AC 之后别急着走试着用其他方法重写一遍对比它们的特点比如二进制枚举写法for (int mask 0; mask (1 n); mask) { if (__builtin_popcount(mask) ! k) continue; int sum 0; for (int i 0; i n; i) { if (mask (1 i)) sum a[i]; } if (isPrime(sum)) ans; }再比如“选/不选”递归写法void dfs(int idx, int cnt, int sum) { if (idx n) { if (cnt k isPrime(sum)) ans; return; } dfs(idx 1, cnt 1, sum a[idx]); // 选当前数 dfs(idx 1, cnt, sum); // 不选当前数 }三种写法搜索的空间规模完全不同。start 参数版是 C(n, k)二进制枚举是 2^n选/不选递归也是 2^n。当 n 20 时差别不大但当 n 到 30 的时候C(30, 15) ≈ 1.55 亿而 2^30 ≈ 10 亿差距就出来了。理解每种写法的复杂度来源以后你才能根据数据范围在第一时间选出最合适的方案。6.3 数据范围变大怎么办如果你做完这道题还觉得不过瘾可以想一想如果 n 变成 40 呢这时候 C(40, 20) 大约是 1378 亿直接 DFS 肯定 TLE。这种情况下就要请出折半搜索meet-in-the-middle把 n 个数分成两半各自枚举出所有选择的和然后想办法合并。具体到这道题就是先枚举左半边选 0 ~ k 个数的所有和再枚举右半边选 0 ~ k 个数的所有和对于左半边的每一种和 left需要在右半边找 k - cnt_left 个数构成的 right使得 left right 是素数。复杂度大约是 O(2^(n/2) * n) 量级能够轻松应对 n ≤ 40 的数据。这就是 P1036 的加强版思路也是我在刷题后期才逐渐掌握的技巧。我个人在实际操作中的体会是做这种基础题最重要的不是背模板而是把“为什么这样写”想清楚。为什么 start 参数能去重为什么试除法用i * i x为什么 sum 用值传递每一个小决定背后都有它的道理。把这些道理吃透比刷一百道类似题都管用。最后再分享一个小技巧AC 之后建议尝试改一改代码把 start 参数去掉看看会得到什么结果把 isPrime 的i * i改成i * i x 1看看会不会出错亲手触碰边界你才能真正理解这道题的设计精妙之处。