Question

我有以下数组：

a = np.arange(12).reshape((2, 2, 3))

和

b = np.zeros((2, 2))

现在我要使用b来访问a，s.t.如果b [i，j] = z，则在每个索引i，j处取a的第z个元素。对于以上示例的含义，答案应为[[0，3]，[6，9]]。我觉得这与np.choose非常相关，但是以某种方式无法完全管理它。你能帮我吗？

Answer 1

可以建议两种方法。

使用advanced-indexing的explicit个范围数组-

m,n = b.shape
out = a[np.arange(m)[:,None],np.arange(n),b.astype(int)]

使用np.take_along_axis-

np.take_along_axis(a,b.astype(int)[...,None],axis=2)[...,0]

样品运行-

In [44]: a
Out[44]: 
array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]])

In [45]: b
Out[45]: 
array([[0., 0.],
       [0., 0.]])

In [46]: m,n = b.shape

In [47]: a[np.arange(m)[:,None],np.arange(n),b.astype(int)]
Out[47]: 
array([[0, 3],
       [6, 9]])

In [48]: np.take_along_axis(a,b.astype(int)[...,None],axis=2)[...,0]
Out[48]: 
array([[0, 3],
       [6, 9]])

使用数组作为多维数组的索引掩码

1 个答案: