如何将嵌套Python生成器转为扁平生成器以适配并行计算？

你当前的核心问题是：nk_pairs里的每个元素是生成器对象+参数k，但pqdm需要的是可独立处理的单个任务单元，而且生成器无法在多进程间序列化传递，导致并行逻辑失效。同时要保持内存高效，不能一次性把所有组合加载到内存。

原代码问题点 #

原代码中nk_pairs是集合，每个元素是(combinations_with_replacement生成器, k)，但pqdm会把每个生成器当作单独任务传递，进程无法正确迭代生成器内容，且生成器不可序列化，这是根本原因。

修改后的完整代码 #

from pqdm.processes import pqdm
from itertools import combinations_with_replacement, chain
from math import pow

max_iter = 5  # 原1+4，直接写5更直观
max_a = 1000

# 生成扁平的任务流：逐个产出(组合数组, k)，内存零负担
def generate_tasks():
    for k in range(1, max_iter):
        for n in range(1, k+1):
            # 遍历每个组合，和对应的k配对成单个任务
            for combo in combinations_with_replacement(range(1, max_a), r=n):
                yield (combo, k)

# 修正整数k次方判断的精度问题：用整数运算替代浮点数
def clear_root(val, k):
    """判断val是否为某个正整数的k次方"""
    root = round(val ** (1/k))
    # 检查相邻三个整数的k次方，避免浮点数近似误差
    for candidate in [root-1, root, root+1]:
        if candidate <= 0:
            continue
        if candidate ** k == val:
            return True
    return False

# 适配pqdm的任务格式：接受单个任务元组
def func(task):
    array, k = task
    # 用整数幂运算替代pow，避免浮点数精度损失
    sum_ = sum(x ** k for x in array)
    if clear_root(sum_, k):
        return array, k

# 并行执行：生成器作为输入，内存高效利用多核
results = pqdm(generate_tasks(), func, n_jobs=24, argument_type='args')

关键修改说明 #

1. 扁平任务生成：用generate_tasks生成器函数，把所有(组合, k)的配对逐个产出，既保持内存高效（不会一次性加载所有组合），又给pqdm提供了可直接处理的单个任务单元。
2. 修复精度问题：原clear_root用浮点数计算k次方根容易出现精度误差，改成先取近似整数，再验证相邻整数的k次方是否等于目标值，完全用整数运算避免误差。
3. 适配并行参数：func接受单个任务元组，配合pqdm的argument_type='args'，自动把元组解包为函数参数（或者也可以直接让func接受两个参数，pqdm会自动传递元组里的元素）。

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