多维度的Numpy花式索引

Let's say I have an numpy array A of size n x m x k and another array B of size n x m that has indices from 1 to k. I want to access each n x m slice of A using the index given at this place in B, giving me an array of size n x m.

假设我有一个大小为n x m x k的numpy数组A和另一个大小为n x m的数组B,其索引从1到k。我想使用B中这个地方给出的索引访问每个n x m切片的A,给我一个大小为n x m的数组。

Edit: that is apparently not what I want! [[ I can achieve this using take like this:

编辑:这显然不是我想要的! [[我可以用这样的方式实现这个目的:

A.take(B) ]] end edit

A.take(B)]]结束编辑

Can this be achieved using fancy indexing? I would have thought A[B] would give the same result, but that results in an array of size n x m x m x k (which I don't really understand).

这可以通过花式索引来实现吗?我原以为A [B]会给出相同的结果,但是会产生一个大小为n x m x m x k的数组(我真的不明白)。

The reason I don't want to use take is that I want to be able to assign this portion something, like

我不想使用take的原因是我希望能够分配这个部分,比如

A[B] = 1

A [B] = 1

The only working solution that I have so far is

到目前为止,唯一可行的解决方案是

A.reshape(-1, k)[np.arange(n * m), B.ravel()].reshape(n, m)

A.reshape(-1,k)[np.arange(n * m),B.ravel()]。reshape(n,m)

but surely there has to be an easier way?

但肯定有一个更简单的方法吗?

1 个解决方案

#1

Suppose

import numpy as np
np.random.seed(0)

n,m,k = 2,3,5
A = np.arange(n*m*k,0,-1).reshape((n,m,k))
print(A)
# [[[30 29 28 27 26]
#   [25 24 23 22 21]
#   [20 19 18 17 16]]

#  [[15 14 13 12 11]
#   [10  9  8  7  6]
#   [ 5  4  3  2  1]]]

B = np.random.randint(k, size=(n,m))
print(B)
# [[4 0 3]
#  [3 3 1]]

To create this array,

要创建此数组,

print(A.reshape(-1, k)[np.arange(n * m), B.ravel()])
# [26 25 17 12  7  4]

as a nxm array using fancy indexing:

作为使用花式索引的nxm数组:

i,j = np.ogrid[0:n, 0:m]
print(A[i, j, B])
# [[26 25 17]
#  [12  7  4]]

#1