如何让itertools.product生成的组合按元素和递增且无需存储全量组合

实现思路 #

不需要预先生成所有笛卡尔积组合存储到内存，我们可以直接按总和从小到大的顺序逐层生成合法组合：

1. 先获取三个range的取值上下限，计算得到总和的最小、最大可能值
2. 从最小总和到最大总和依次遍历，对每个当前总和，仅枚举能凑出该总和的合法元素组合，直接返回即可，全程不会存储所有组合数据

实现代码 #

def product_ordered_by_sum(a_range, b_range, c_range):
    # 提取三个range的取值边界
    a_min, a_max = a_range.start, a_range.stop - 1
    b_min, b_max = b_range.start, b_range.stop - 1
    c_min, c_max = c_range.start, c_range.stop - 1

    # 计算总和的取值范围
    s_min = a_min + b_min + c_min
    s_max = a_max + b_max + c_max

    # 按总和从小到大遍历
    for s in range(s_min, s_max + 1):
        # 裁剪a的合法取值范围，减少无效遍历
        a_low = max(a_min, s - b_max - c_max)
        a_high = min(a_max, s - b_min - c_min)
        for a_val in range(a_low, a_high + 1):
            sum_bc = s - a_val
            # 裁剪b的合法取值范围
            b_low = max(b_min, sum_bc - c_max)
            b_high = min(b_max, sum_bc - c_min)
            for b_val in range(b_low, b_high + 1):
                c_val = sum_bc - b_val
                yield (a_val, b_val, c_val)

调用示例 #

import numpy as np

a = np.arange(0,5)
b = np.arange(0,3)
c = np.arange(0,7)

# 直接迭代生成器，内存仅占用当前生成的单个组合
for p in product_ordered_by_sum(a, b, c):
    print(p)

输出效果 #

和要求的顺序完全一致，总和相同的组合不需要额外排序，符合需求：

(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(1, 0, 0)
(0, 1, 1)
(1, 0, 1)
(1, 1, 0)
(0, 0, 2)
...

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