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]

我尝试过以下方法来获得我想要的东西:

ds = np.array([[0,3],[0,3],[0,10]])
nItems = int(reduce(lambda a,b: a * (b[1] - b[0] + 1), ds, 1))
myCombinations = np.zeros((nItems,))
nArrays = []
for x in range(ds.shape[0]):
    low = ds[x][0]
    high= ds[x][1]
    nitm = high - low + 1
    ar = [x+low for x in range(nitm) ]
    nArrays.append(ar)

myCombinations = cartesian(nArrays)

笛卡尔函数取自于使用numpy来构建两个数组的所有组合的数组

我需要做几百万次.

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

Answer 1

我想你要找的是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

说明

np.mgrid在N维空间中为您提供一组网格点.让我试着用一个更小的例子来说明这一点,以使事情变得更清楚:

>>> a = np.mgrid[0:2, 0:2]
>>> a
array([[[0, 0],
        [1, 1]],

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

由于我给出了两组范围0:2, 0:2,我得到了一个2D网格.什么mgrid返回是x的值和对应于该网格点(0,0),(0,1),(1,0)和(1,1)在2D空间中的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轴放在最后一个轴之后的方法."你可能也想读这个.

>>> 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.