我有两个numpy的阵列,A
和indices
。numpy的匹配索引尺寸
A
具有尺寸m x n x 10000. indices
具有尺寸m x n x 5(从argpartition(A, 5)[:,:,:5]
输出)。 我想得到一个m×n×5的数组,其中包含对应于indices
的A
的元素。
尝试
indices = np.array([[[5,4,3,2,1],[1,1,1,1,1],[1,1,1,1,1]],
[500,400,300,200,100],[100,100,100,100,100],[100,100,100,100,100]])
A = np.reshape(range(2 * 3 * 10000), (2,3,10000))
A[...,indices] # gives an array of size (2,3,2,3,5). I want a subset of these values
np.take(A, indices) # shape is right, but it flattens the array first
np.choose(indices, A) # fails because of shape mismatch.
动机
我试图得到A[i,j]
5个最大值为每i<m
,j<n
使用np.argpartition
因为阵列可以得到相当大的排序顺序。