python - itertools 产品加速

Question

我使用 itertools.product 生成长度为 13 的 4 个元素的所有可能变体。4 和 13 可以是任意的，但实际上，我得到 4^13 个结果，这很多。我需要将结果作为 Numpy 数组，目前执行以下操作:

  c = it.product([1,-1,np.complex(0,1), np.complex(0,-1)], repeat=length)
  sendbuf = np.array(list(c))

中间插入了一些简单的分析代码，第一行看起来几乎是即时的，而转换为列表和 Numpy 数组大约需要 3 个小时。有没有办法让它更快？这可能是我忽略的非常明显的事情。

谢谢!

Answer 1

itertools.product() 的 NumPy 等效项是 numpy.indices()，但它只会为您提供 0,.. .,k-1:

numpy.rollaxis(numpy.indices((2, 3, 3)), 0, 4)
array([[[[0, 0, 0],
         [0, 0, 1],
         [0, 0, 2]],

        [[0, 1, 0],
         [0, 1, 1],
         [0, 1, 2]],

        [[0, 2, 0],
         [0, 2, 1],
         [0, 2, 2]]],


       [[[1, 0, 0],
         [1, 0, 1],
         [1, 0, 2]],

        [[1, 1, 0],
         [1, 1, 1],
         [1, 1, 2]],

        [[1, 2, 0],
         [1, 2, 1],
         [1, 2, 2]]]])

对于您的特殊情况，您可以使用

a = numpy.indices((4,)*13)
b = 1j ** numpy.rollaxis(a, 0, 14)

(这不会在 32 位系统上运行，因为数组太大。根据我可以测试的大小推断，它应该在不到一分钟的时间内运行。)

EIDT:顺便提一下:对 numpy.rollaxis() 的调用或多或少是装饰性的，以获得与 itertools.product() 相同的输出。如果你不关心索引的顺序，你可以忽略它(但只要你没有任何后续操作会将你的数组转换为连续数组，它就很便宜。)

EDIT2:得到的精确模拟

numpy.array(list(itertools.product(some_list, repeat=some_length)))

你可以使用

numpy.array(some_list)[numpy.rollaxis(
    numpy.indices((len(some_list),) * some_length), 0, some_length + 1)
    .reshape(-1, some_length)]

这变得完全不可读——请告诉我是否应该进一步解释它:)