加速 itertools product

Question

我使用itertools.product来生成4个元素长度为13的所有可能组合。这里的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

你可以通过跳过转换成列表的步骤来加快速度：

numpy.fromiter(c, count=…)  # Using count also speeds things up, but it's optional

这个函数首先分配了一个NumPy数组，然后逐个初始化每个元素，而不需要额外的步骤去构建一个列表。

附注：fromiter() 不能处理 product() 返回的元组，所以现在可能不是一个解决方案。如果 fromiter() 能处理 dtype=object，那就可以用了。

再附注：正如Joe Kington指出的，可以通过将元组放入结构化数组来让这个方法有效。不过，这样做似乎并不总是能加快速度。

Answer 2

第一行看起来是瞬间完成的，因为实际上并没有进行任何操作。生成器对象只是被创建了，只有当你开始遍历它时，才会真正进行操作。正如你所说，你得到了 4^13 = 67108864 个数字，这些数字都是在你调用 list 的时候计算出来并准备好的。我注意到 np.array 只接受列表或元组，所以你可以尝试把你的迭代器转换成一个元组，然后传给 np.array，看看这样做是否会有性能上的差异，并且不会影响你程序的整体性能。不过，这个只能通过尝试你的具体情况来确定，虽然有一些 观点 表示元组会稍微快一点。

如果想用元组来试试，直接用元组替代列表，操作如下：

sendbuf = np.array(tuple(c))

Answer 3

NumPy中与 itertools.product() 相似的功能是 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位系统上运行不了，因为数组太大。不过根据我能测试的大小，它应该在一分钟内完成。）

补充说明：提到 numpy.rollaxis() 这个调用其实主要是为了让输出和 itertools.product() 一样。如果你不在意索引的顺序，可以直接省略这个步骤（不过只要后面没有其他操作会把你的数组变成连续的数组，这个操作也不会消耗太多资源）。

补充说明2：要得到和

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)]

这段代码变得完全看不懂了——只要告诉我是否需要进一步解释一下 :)