如何高效遍历嵌套数组查找元素和为40的组合，解决超大规模样本运算慢问题

原来的实现逻辑是先全量生成7组长度为16的数组的所有笛卡尔积（总规模16^7=268,435,456约2.7亿），再逐行校验和是否等于40，无论是内存占用还是遍历计算的开销都非常大，必然会出现耗时极长的问题。

我们采用折半枚举法做优化，将7个数组拆分为两组分别计算，计算规模直接下降几个数量级：

• 前3个数组为第一组，总组合数为16^3=4096
• 后4个数组为第二组，总组合数为16^4=65536
  两组总枚举量不到7万，仅为原方案的0.026%，运行耗时不到1秒即可得到全部结果。

优化后代码 #

import numpy as np
from collections import defaultdict

def get_partial_combos(arrays):
    # 计算分组的所有组合和对应的组合值
    la = arrays.shape[0]
    grids = np.ix_(*arrays)
    combos = np.stack(grids, axis=-1).reshape(-1, la)
    sum_vals = combos.sum(axis=1)
    return sum_vals, combos

rows = np.array(
    [
        [2, 15, 23, 19, 3, 2, 3, 27, 20, 11, 27, 10, 19, 10, 13, 10],
        [22, 9, 5, 10, 5, 1, 24, 2, 10, 9, 7, 3, 12, 24, 10, 9],
        [16, 0, 17, 0, 2, 0, 2, 0, 10, 0, 15, 0, 6, 0, 9, 0],
        [11, 27, 14, 5, 5, 7, 8, 24, 8, 3, 6, 15, 22, 6, 1, 1],
        [10, 0, 2, 0, 22, 0, 2, 0, 17, 0, 15, 0, 14, 0, 5, 0],
        [1, 6, 10, 6, 10, 2, 6, 10, 4, 1, 5, 5, 4, 8, 6, 3],
        [6, 0, 13, 0, 3, 0, 3, 0, 6, 0, 10, 0, 10, 0, 10, 0],
    ],
    dtype="int32",
)

# 拆分数组为两部分
split_pos = 3
part1_sum, part1_combo = get_partial_combos(rows[:split_pos])
part2_sum, part2_combo = get_partial_combos(rows[split_pos:])

# 构建第二部分的和到组合的映射索引
sum2combos = defaultdict(list)
for s, combo in zip(part2_sum, part2_combo):
    sum2combos[s].append(combo)

# 匹配符合和为40的组合
combos = []
target = 40
for s1, c1 in zip(part1_sum, part1_combo):
    need_s2 = target - s1
    if need_s2 in sum2combos:
        for c2 in sum2combos[need_s2]:
            combos.append(np.concatenate([c1, c2]))

# 输出结果
print(f"共找到{len(combos)}个符合条件的组合：")
for combo in combos:
    print(combo)

核心优化点 #

• 避免全量笛卡尔积生成，通过分组枚举将计算量降到可忽略的量级
• 用哈希映射存储第二组的和与组合对应关系，匹配阶段直接O(1)查找补数，无需二次遍历

内容的提问来源于stack exchange，提问作者mhay10