如何让itertools组合均匀“增加”？

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 之外的步骤，但是对于越来越大的步骤，这样做的效率会降低，因为最后一个范围可能不包含生成当前总和所需的元素。这似乎是不可避免的，因为此时您需要解决一个困难的计算问题，才能按总和有效地循环这些乘积。