问题

我有两个numpy数组,A和indices.

A尺寸为mxnx 10000. indices尺寸为mxnx 5(输出自argpartition(A, 5)[:,:,:5]).我想获得amxnx 5阵列含有的元素A对应于indices.

尝试

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<nnp.argpartition