python - 如何创建包含数字范围的矩阵？

Question

我想创建包含整数范围的矩阵，例如

[[range(1, 3), 0, 0],
 [range(5, 7), range(0, 2), 0],
 [0, range(1, 3), 0]]

我期望的答案是

[[[1 0 0], [5 0 0], [0, 1, 0]]
 [[2 0 0], [5 0 0], [0, 1, 0]]
 [[1 0 0], [6 0 0], [0, 1, 0]]
 [[2 0 0], [6 0 0], [0, 1, 0]]
 [[1 0 0], [5 0 0], [0, 2, 0]]
 [[2 0 0], [5 0 0], [0, 2, 0]]
 [[1 0 0], [6 0 0], [0, 2, 0]]
 [[2 0 0], [6 0 0], [0, 2, 0]]]

而不是矩阵内的范围对象。

我可以通过使用循环来实现我的目标

def foo(p1_range, p2_range, p3_range)
    for p1 in range(p1_range):
        for p2 in range(p2_range):
            for p3 in range(p3_range):
                yield [[p1 0 0], [p2 0 0], [0, p3, 0]]

有什么pythonic的方法可以解决这个问题吗？

注意:我的主代码通常需要 7 个参数...这种方式看起来很糟糕。

Answer 1

您可以使用 itertools.product减少缩进级别并消除嵌套循环:

from itertools import product

def foo(p1_range, p2_range, p3_range):
    for p1, p2, p3 in product(p1_range, p2_range, p3_range):
        yield [[p1, 0, 0], [p2, 0, 0], [0, p3, 0]]