Question

我一直在努力寻找这个问题的答案，但没有找到任何非常适合这个问题的答案。我有一个多维的numpy数组包含数据（在我的情况下是3维）和另一个数组（2维），其中包含我想要沿原始数组的最后一个维度的值的信息。例如，这是一个说明问题的简单示例。我有一个数组a，另一个数组b包含a维2的索引。我想要一个新的二维数组c，其中c[i, j] = a[i, j, b[i, j]]。我能想到的唯一方法是使用循环，如下所述。然而，这似乎很笨拙和缓慢。

In [3]: a = np.arange(8).reshape((2, 2, 2))
In [4]: a
Out[4]: 
array([[[0, 1],
        [2, 3]],

       [[4, 5],
        [6, 7]]])

In [6]: b = np.array([[0, 1], [1, 1]])

In [8]: c = np.zeros_like(b)

In [9]: for i in xrange(2):
   ...:     for j in xrange(2):
   ...:         c[i, j] = a[i, j, b[i, j]]

In [10]: c
Out[10]: 
array([[0, 3],
       [5, 7]])

有没有更多的pythonic方式来做这个，也许是我不知道的一些numpy索引功能？

Answer 1

当您想要使用多维数组索引多维数组时，每个维度的索引将一起广播。考虑到这一点，你可以这样做：

>>> rows = np.arange(a.shape[0])
>>> cols = np.arange(a.shape[1])
>>> a[rows[:, None], cols, b]
array([[0, 3],
       [5, 7]])

Answer 2

In [40]: a = np.arange(8).reshape((2, 2, 2))

In [41]: b = np.array([[0, 1], [1, 1]])

In [42]: i = np.array([[0,0],[1,1]])

In [43]: a[i,i.T,b]
Out[43]: 
array([[0, 3],
       [5, 7]])

或使用ix_生成索引：

In [47]: j = np.ix_([0,1],[0,1])

In [48]: a[j[0],j[1],b]
Out[48]: 
array([[0, 3],
       [5, 7]])
In [49]: j
Out[49]: 
(array([[0],
       [1]]), array([[0, 1]]))

或ogrid

In [101]: i = np.ogrid[0:2,0:2]

In [102]: i.append(b)

In [103]: a[i]
Out[103]: 
array([[0, 3],
       [5, 7]])

Numpy多维子集

2 个答案: