使用2d数组作为n-D数组的索引列表

时间:2016-03-08 17:09:46

标签: arrays numpy multidimensional-array indexing

如果我有以下数据:

A = np.random.random((3, 4, 5))

# np.all(indices < A.shape) is true
indices = np.array([
    [0, 0, 0],
    [1, 2, 4],
    ...
    [2, 3, 4]
])

如何将indices的每一行作为一组轴索引用于A以提供以下内容?

B = np.array([
    A[0, 0, 0],
    A[1, 2, 4],
    ...
    A[2, 3, 4]
])

1 个答案:

答案 0 :(得分:2)

这是一个第二个例子:

In [1]: A=np.arange(10,22).reshape(3,4)
In [2]: A
Out[2]: 
array([[10, 11, 12, 13],
       [14, 15, 16, 17],
       [18, 19, 20, 21]])
In [3]: ind=np.array([[0,1],[1,3],[2,0],[0,2]])
In [4]: ind
Out[4]: 
array([[0, 1],
       [1, 3],
       [2, 0],
       [0, 2]])
In [5]: A[ind[:,0],ind[:,1]]
Out[5]: array([11, 17, 18, 12])

或您的变量,

A[indices[:,0], indices[:,1], indices[:,2]]

或更一般地说:

In [8]: tuple(ind.T)
Out[8]: (array([0, 1, 2, 0]), array([1, 3, 0, 2]))
In [9]: A[tuple(ind.T)]
Out[9]: array([11, 17, 18, 12])

这是基于A[a,b]A[(a,b)]相同的想法。当ab匹配列表或数组时,它会通过配对来选择值,大致与

相同
[A[i,j] for i,j in zip(a,b)]

对于类似product的索引,索引数组需要有更多维度。 ix_是生成此类数组的便捷方法:

In [53]: np.ix_(ind[:,0],ind[:,1])
Out[53]: 
(array([[0],
        [1],
        [2],
        [0]]), array([[1, 3, 0, 2]]))
In [54]: A[np.ix_(ind[:,0],ind[:,1])]
Out[54]: 
array([[11, 13, 10, 12],
       [15, 17, 14, 16],
       [19, 21, 18, 20],
       [11, 13, 10, 12]])

In [56]: A[ind[:,[0]],ind[:,1]]
Out[56]: 
array([[11, 13, 10, 12],
       [15, 17, 14, 16],
       [19, 21, 18, 20],
       [11, 13, 10, 12]])