python - 如何使 itertools 组合 'increase' 均匀？

Question

考虑以下示例:

import itertools
import numpy as np

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

prods = itertools.product(a,b,c)

for p in prods:
    print(p)

这按以下顺序迭代产品:

(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 0, 3)
(0, 0, 4)
(0, 1, 0)

但我更愿意按照总和的顺序给出产品，例如

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

如何在不将所有组合存储在内存中的情况下实现这一目标？
注: a b和 c总是范围，但不一定具有相同的最大值。当两个乘积的和相等时，也没有二级排序，即 (0,1,1)相当于 (2,0,0) .

Answer 1

无需在内存中存储额外产品，最简单的方法是使用递归。而不是 range(a,b) ，传入(a,b)的列表配对并自己进行迭代:

def prod_by_sum(range_bounds: List[Tuple[int, int]]):
    """
    Yield from the Cartesian product of input ranges, produced in order of sum.

    >>> range_bounds = [(2, 4), (3, 6), (0, 2)]
    >>> for prod in prod_by_sum(range_bounds):
    ...    print(prod)
    (2, 3, 0)
    (2, 3, 1)
    (2, 4, 0)
    (3, 3, 0)
    (2, 4, 1)
    (2, 5, 0)
    (3, 3, 1)
    (3, 4, 0)
    (2, 5, 1)
    (3, 4, 1)
    (3, 5, 0)
    (3, 5, 1)

    """
    def prod_by_sum_helper(start: int, goal_sum: int):
        low, high = range_bounds[start]
        if start == len(range_bounds) - 1:
            if low <= goal_sum < high:
                yield (goal_sum,)
            return

        for current in range(low, min(high, goal_sum + 1)):
            yield from ((current,) + extra
                        for extra in prod_by_sum_helper(start + 1, goal_sum - current))

    lowest_sum = sum(lo for lo, hi in range_bounds)
    highest_sum = sum(hi - 1 for lo, hi in range_bounds)

    for goal_sum in range(lowest_sum, highest_sum + 1):
        yield from prod_by_sum_helper(0, goal_sum)

其中有输出 range_bounds = [(0, 5), (0, 3), (0, 7)]从...开始:

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

您可以通过修改单个列表并生成它的副本来迭代地执行此精确过程，但是代码要么变得更复杂，要么效率更低。
您也可以简单地修改它以支持除 1 之外的步骤，但是随着越来越大的步骤，它的工作效率会降低，因为最后一个范围可能不包含生成当前总和所需的元素。这似乎是不可避免的，因为在这一点上，您需要解决一个困难的计算问题才能有效地按总和循环遍历这些乘积。