B002排序、双指针、哈希表、两数之和到K数之和（1640~1642）CSES？

1640 Sum of Two Values - CSES 1641 Sum of Three Values - CSES 1642 Sum of Four Values - CSES 这类问题可以抽象成 \(K\) 数之和，本质是递归的转化为两数之和，时间复杂度为 \(N^{K-1}\) 。一般可以排序双指针或哈希表解决。 题意：给定一个 \(n\) 个整数的数组，找出 \(k\) 个不同位置（1-based）的元素，使其的和等于 \(x\) 。若无解输出 IMPOSSIBLE 。与力扣上的题目的有差异。 两数之和 因为需要输出位置，所以将原数组变成一个元组数组再排序。 排完序后数组具有单调性，前后两个指针的和与 \(x\) 比大小，小就 l ++ 大就 r -- 。 具体代码为 import sys from math import inf if 1: inp = lambda: sys.stdin.readline().strip() II = lambda: int(inp()) MII = lambda: map(int, inp().split()) LII = lambda: list(MII()) Max = lambda x, y: x if x > y else y Min = lambda x, y: x if x < y else y def main(): n, k = MII() a = LII() a = [(x, i+1) for i, x in enumerate(a)] a.sort() l, r = 0, n - 1 while l < r: s = a[l][0] + a[r][0] if s > k: r -= 1 elif s < k: l += 1 else: print(a[l][1], a[r][1]) return print("IMPOSSIBLE") if __name__ == "__main__": main() 还有更有学习意义的算法，即使用哈希表边枚举边查询。这个技巧很有用，很多场景用到的有很大的启发意义。 # 模版代码直接复制即可 def main(): n, k = MII() a = LII() mp = {} for i, x in enumerate(a): if k - x in mp: # 查是否遍历过 print(mp[k - x], i + 1) return mp[x] = i + 1 # 记录出遍历的部分 print("IMPOSSIBLE") 三数之和 在两数之和的外面套一层 for 循环转化求剩余区间的两数之和。 具体代码为 # 模版代码直接复制即可 def main(): n, k = MII() a = LII() a = [(x, i+1) for i, x in enumerate(a)] a.sort() for i in range(n): l, r = i + 1, n - 1 while l < r: s = a[l][0] + a[r][0] + a[i][0] if s > k: r -= 1 elif s < k: l += 1 else: print(a[i][1], a[l][1], a[r][1]) return print("IMPOSSIBLE") 四数之和 按照三数之和的思路套两层 for 循环时会超时，这时考虑剪枝和哈希表。 直觉告诉我们可以预先估计数的大小然后及时止损。 比如 2, 3, 5, 8, 11 和 x = 8 ，当 \(a_i=3\) 的时候已经没有往后遍历的必要了，因为有序数组在 \(i\) 以后的元素都大于等于 \(a_i\) ，所以这里最小也是 \(9\) ，不可能由等于 \(8\) 的机会了。 于是我们可以在两个循环的开始进行预测一下有没有答案存在的可能。 for i in range(n): if a[i][0] * 4 > k: break for j in range(i + 1, n): if a[i][0] + a[j][0] * 3 > k: break # while l < n: # ... 这个剪枝是很粗糙的，勉强可以过，但是我们还可以再精益求精一下。 最小值剪枝：如果当前数 + 后面三个都大于 \(K\) ，说明往后都不存在了直接 break 最大值剪枝：如果当前数 + 后面三个都小于 \(K\) ，说明往后都不存在了 continue 具体代码为 for i in range(n - 3): # 剪枝 1 if a[i][0] + a[i+1][0] + a[i+2][0] + a[i+3][0] > k: break # 剪枝 2 if a[i][0] + a[n-1][0] + a[n-2][0] + a[n-3][0] < k: continue for j in range(i + 1, n - 2): # 二级剪枝 if a[i][0] + a[j][0] + a[j+1][0] + a[j+2][0] > k: break if a[i][0] + a[j][0] + a[n-1][0] + a[n-2][0] < k: continue # while l < r: # ... 使用哈希表转化为两对数的和为 \(K\) 思路与两数之和相似关键是去重。 将 \(A+B+C+D=K\) 转化为 \((A+B)=K-(C+D)\) 。按照两数之和的思路进行枚举和查询。 最外层枚举 \(i\) 作为分割点 内层第一个循环 \(j\) 从 \(i+1\) 开始查询 第二个循环 \(j\) 从 \(0\) 开始枚举到 \(i-1\) ，插入到哈希表中。这样保证不会下标重复。 # 模版代码直接复制即可 def main(): n, k = MII() a = LII() a = [x for x in a] mp = {} for i in range(n): for j in range(i + 1, n): t = k - a[i] - a[j] if t in mp: p, q = mp[t] print(i + 1, j + 1, p, q) return for j in range(i): # 防重，只加入以i为右边界的组合，与只枚举以i右边界的对应 mp[a[i] + a[j]] = (i + 1, j + 1) print("IMPOSSIBLE") K数之和 递归和剪枝，学习中，模版由 Gemini 3.1 Pro 给出 lol 模版代码为 import sys def main(): # 快速 I/O input = sys.stdin.read data = input().split() if not data: return n = int(data[0]) target = int(data[1]) # 填入你需要求的 K (2, 3, 4...) # 比如 CSES 1641 就填 3，CSES 1642 就填 4 K = 4 arr = [int(x) for x in data[2:]] # 绑定原数组的 1-based 下标，并按数值排序 # a[i] = (数值, 原下标) a =[(val, i + 1) for i, val in enumerate(arr)] a.sort() # 递归函数：在 a[start:] 中寻找 k 个数，使其数值之和为 target # 返回值：一个包含 k 个原数组下标的列表，若找不到则返回 None def k_sum(start, target, k): # 边界保护 if k < 2 or start + k > n: return None # 【基础情况】：当 k 降维到 2 时，使用双指针 O(N) 夹逼 if k == 2: l, r = start, n - 1 while l < r: s = a[l][0] + a[r][0] if s == target: return [a[l][1], a[r][1]] elif s < target: l += 1 else: r -= 1 return None # 【递归情况】：当 k > 2 时，固定一个数，将 k 降维为 k-1 for i in range(start, n - k + 1): # 剪枝 1：去重优化。如果有连续相同的数值，只把第一个作为该层循环的起点，避免重复计算 if i > start and a[i][0] == a[i - 1][0]: continue # 剪枝 2：当前数值与能组成的【最小】的 k-1 个数之和都已经大于 target， # 说明后面的数更不可能，直接跳出循环！ if a[i][0] * k > target: break # 剪枝 3：当前数值与能组成的【最大】的 k-1 个数之和都小于 target， # 说明当前数太小了，没资格参与，跳过进入下一个数！ if a[i][0] + a[-1][0] * (k - 1) < target: continue # 递归调用降维寻找 sub_res = k_sum(i + 1, target - a[i][0], k - 1) # 如果找到了答案，立刻拼接当前数的下标并返回 if sub_res: return [a[i][1]] + sub_res return None # 调用递归 ans = k_sum(0, target, K) # 按格式输出结果 if ans: print(*(ans)) # 输出由空格分隔的下标 else: print("IMPOSSIBLE") if __name__ == "__main__": main()