如何加速Python中itertools.combinations处理54选16组合的运算？

前置说明 #

54选16的全量组合共有约2.1×10^13组，即便优化到每1纳秒处理1组，也需要连续运行230天才能跑完，全量运算不存在可行性。以下优化方案针对处理1%量级组合的需求设计，可将运行耗时降低10~20倍。

现有代码的核心冗余问题 #

• 对生成的组合转set：itertools.combinations生成的组合本身无重复，转set会触发全量遍历哈希化，浪费大量算力
• 逐个拆分元素赋值：手动给p1-p16赋值再建数组的操作完全多余，可通过numpy的reshape方法一步完成
• 无用操作：vun列表存储打乱后的组合、sample(i, num2)随机打乱操作如果后续无使用，可直接删除，节省随机数生成开销
• 单线程运行：i7处理器多为8核16线程，单线程仅能利用不到10%的CPU算力
• 全量存储数组：如果将所有4x4数组存入tabl列表，处理1%组合时需要至少3TB内存，完全不现实，需改为边生成边计算，不存储中间数组

优化后代码示例 #

import numpy as np
import itertools
from itertools import islice
from multiprocessing import Pool, cpu_count

# 可调整参数
num1 = 54
num2 = 16
sample_rate = 0.01  # 处理1%的组合
process_num = cpu_count()  # 用满所有CPU核心

# 提前生成数字列表
nums = list(range(1, num1+1))
total_combin = int(itertools.combinations(range(num1), num2).__length_hint__())
step = int(1 / sample_rate)
total_task = total_combin // step

# 单进程处理函数：处理指定分片的组合
def process_chunk(chunk_idx):
    chunk_size = total_task // process_num
    start = chunk_idx * chunk_size * step
    end = start + chunk_size * step if chunk_idx != process_num-1 else total_combin
    counter = 0
    # 可在这里替换为你自己的对4x4数组的计算逻辑
    for comb in islice(itertools.combinations(nums, num2), start, end, step):
        arr = np.array(comb, dtype=np.int8).reshape(4,4)
        # 这里加你对arr的处理逻辑，不要存储arr到全局列表
        counter += 1
    return counter

if __name__ == "__main__":
    # 多进程并行处理
    with Pool(process_num) as pool:
        counters = pool.map(process_chunk, range(process_num))
    print("总处理组合数：", sum(counters))

额外优化建议 #

• 如果对组合的处理逻辑可以向量化，可每次读取1000~10000组组合，一次性转为(N, 4, 4)的三维numpy数组批量计算，可再提升3~5倍性能
• 如果不需要固定步长采样，而是需要随机采样1%的组合，可使用numpy.random.choice提前生成要采样的组合索引，再通过itertools.islice定位到对应组合，避免遍历不需要的组合

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