Numpy：从给定范围生成组合的有效方法

Question

我有一个n维数组，如下所示：

np.array([[0,3],[0,3],[0,10]])

在这个数组中，元素表示低值和高值。例如：[0,3]指[0,1,2,3]

我需要使用上面给出的范围生成所有值的组合。 例如，我想要[0,0,0], [0,0,1] ... [0,1,0] ... [3,3,10]

为了得到我想要的东西，我尝试了以下方法：

^{pr2}$

笛卡尔函数取自Using numpy to build an array of all combinations of two arrays

我需要做几百万次。在

我的问题是：有没有更好/有效的方法来做到这一点？在

Answer 1

您可以使用^{}：

In [16]: from itertools import product

In [17]: values = list(product(range(4), range(4), range(11)))

In [18]: values[:5]
Out[18]: [(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 0, 4)]

In [19]: values[-5:]
Out[19]: [(3, 3, 6), (3, 3, 7), (3, 3, 8), (3, 3, 9), (3, 3, 10)]

给定范围数组，可以执行以下操作。（我使用了几个非零低值来演示一般情况，并缩小了输出的大小。：）

^{pr2}$

如果所有范围的低值均为0，则可以使用^{}：

In [52]: values = list(np.ndindex(4, 4, 11))

In [53]: values[:5]
Out[53]: [(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 0, 4)]

In [54]: values[-5:]
Out[34]: [(3, 3, 6), (3, 3, 7), (3, 3, 8), (3, 3, 9), (3, 3, 10)]

Answer 2

我想你要找的是np.mgrid。不幸的是，这将以不同于您需要的格式返回数组，因此您需要进行一些后期处理：

a = np.mgrid[0:4, 0:4, 0:11]     # All points in a 3D grid within the given ranges
a = np.rollaxis(a, 0, 4)         # Make the 0th axis into the last axis
a = a.reshape((4 * 4 * 11, 3))   # Now you can safely reshape while preserving order

说明

^{}给出N维空间中的一组网格点。让我试着用一个小例子来说明这一点，以使事情更清楚：

^{pr2}$

因为我给出了两组范围，0:2, 0:2，所以我得到了一个2D网格。mgrid返回的是二维空间中网格点（0，0），（0，1），（1，0）和（1，1）对应的x值和y值。a[0]告诉你这四个点的x值是什么，a[1]告诉你y值是什么。在

但你真正想要的是我写出来的实际网格点的列表，而不是这些点的x值和y值。第一反应就是根据需要重塑数组：

>>> a.reshape((4, 2))
array([[0, 0],
       [1, 1],
       [0, 1],
       [0, 1]])

但显然这行不通，因为它有效地重塑了扁平数组（通过按顺序读取所有元素得到的数组），而这不是您想要的。在

您要做的是向下查看a的第三个维度，然后创建一个数组：

[ [a[0][0, 0], a[1][0, 0]],
  [a[0][0, 1], a[1][0, 1]],
  [a[0][1, 0], a[1][1, 0]],
  [a[0][1, 1], a[1][1, 1]] ]

上面写着“首先告诉我第一个点（x1，y1），然后是第二个点（x2，y2），…”等等。或许，用一个数字来解释这一点更好。这是a的样子：

                you want to read
                in this direction
                 (0, 0)   (0, 1)
                   |        |
                   |        |
                   v        v

          /        0    0            +  > axis0
 x-values |       /|       /|           /|
          |      / |      / |    axis1 / |
          \     1    1  |         L  |
                |  |     |  |            v
          /     |  0  -| 1           axis2
 y-values |     | /      | /
          |     |/       |/
          \     0    1

                |        |
                |        |
                v        v
              (1, 0)   (1, 1)

np.rollaxis为您提供了一种方法。np.rollaxis(a, 0, 3)在上面的例子中说“取第0个（或最外层的）轴，并使其成为最后一个（或最里面的）轴。（注意：这里实际上只存在轴0、1和2。所以说“将第0个轴发送到第3个位置”是告诉python将第0个轴放在最后一个轴之后的一种方式）。您可能还想阅读this。在

>>> a = np.rollaxis(a, 0, 3)
>>> a
array([[[0, 0],
        [0, 1]],

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

这看起来像你想要的，只是有一个额外的数组维度。我们希望合并维度0和1，以获得一个网格点数组。但是现在平坦的数组以您期望的方式读取，您可以安全地重塑它以获得所需的结果。在

>>> a = a.reshape((4, 2))
>>> a
array([[0, 0],
       [0, 1],
       [1, 0],
       [1, 1]])

3D版也做了同样的事情，只是，我没法画出一个数字，因为它是4D的