python - 为 itertools 重复一个 numpy 数组指定的次数

Question

我正在尝试编写一些代码，以针对不同数量的输入为我提供 itertools 产品。例如，这对我有用。

test = np.array([x for x in itertools.product([0,2],[0,2],[0,2])])

这给了我想要的结果:

>>> test
array([[0, 0, 0],
       [0, 0, 2],
       [0, 2, 0],
       [0, 2, 2],
       [2, 0, 0],
       [2, 0, 2],
       [2, 2, 0],
       [2, 2, 2]])

但是，我希望能够将数量不等的列表传递给产品功能。例如:

test = np.array([x for x in itertools.product([0,2],[0,2],[0,2],[0,2])])

或

test = np.array([x for x in itertools.product([0,2],[0,2])])

我试过了

test = np.array([x for x in itertools.product(([0,2],) * 3)])

和

test = np.array([x for x in itertools.product([[0,2]]*3)])

但都没有给我想要的结果。当然有一个简单的方法可以做到这一点。如果有任何帮助，我将不胜感激。

Answer 1

在我看来，您似乎掌握了 splat-unpack 语法:

>>> n = 3
>>> L = [0, 2]
>>> np.array([x for x in itertools.product(*([L] * n))])
array([[0, 0, 0],
       [0, 0, 2],
       [0, 2, 0],
       [0, 2, 2],
       [2, 0, 0],
       [2, 0, 2],
       [2, 2, 0],
       [2, 2, 2]])

虽然将第二个参数 repeat 用于 itertools.product 可能更容易。

>>> np.array(list(itertools.product(L, repeat=3)))
array([[0, 0, 0],
       [0, 0, 2],
       [0, 2, 0],
       [0, 2, 2],
       [2, 0, 0],
       [2, 0, 2],
       [2, 2, 0],
       [2, 2, 2]])