多个维度的花式索引
问题内容:
假设我有一个大小为nxmxk的numpy数组A和另一个大小为nxm的数组B,其索引从1到k。我想使用在B中此位置给出的索引访问A的每个nxm切片,从而给我一个大小为nx
m的数组。
编辑:那显然不是我想要的![[我可以这样实现take
:
A.take(B)
]]结束编辑
可以使用花式索引来实现吗?我本来以为A[B]
会得到相同的结果,但是会导致数组大小为nxmxmxk(我不太了解)。
我不想使用的原因take
是我希望能够为这部分分配一些内容,例如
A[B] = 1
到目前为止,我唯一可行的解决方案是
A.reshape(-1, k)[np.arange(n * m), B.ravel()].reshape(n, m)
但是肯定有一种更简单的方法吗?
问题答案:
假设
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]]
要创建此数组,
print(A.reshape(-1, k)[np.arange(n * m), B.ravel()])
# [26 25 17 12 7 4]
作为nxm
使用花式索引的数组:
i,j = np.ogrid[0:n, 0:m]
print(A[i, j, B])
# [[26 25 17]
# [12 7 4]]