numpy - numpy：干净地检索最高 k 值的坐标（索引） - 沿特定轴 - 在 ndarray

Question

我希望能够：

1. 沿着（或跨越？）第一个维度选择 k 个最高值
2. 查找这些 k 值的索引
3. 将这些值分配给它们各自位置的形状相同的新 ndarray。

我想知道是否有更快的方法来实现下面举例说明的结果。特别是，我想避免“手动”制作批处理索引。

这是我的解决方案：

# Create unordered array (instrumental to the example)
arr = np.arange(24).reshape(2, 3, 4)
arr_1 = arr[0,::2].copy()
arr_2 = arr[1,1::].copy()
arr[0,::2] = arr_2[:,::-1]
arr[1,1:] = arr_1[:,::-1]

# reshape array to: (batch_size, H*W)
arr_batched = arr.reshape(arr.shape[0], -1)

# find indices for k greatest values along all but the 1st dimension.
gr_ind = np.argpartition(arr_batched, -k)[:, -k]

# flatten and unravel indices.
maxk_ind_flat = gr_ind.flatten()
maxk_ind_shape = np.unravel_index(maxk_ind_flat, arr.shape)
# maxk_ind_shape prints: (array([0, 0, 0, 0]), array([2, 2, 0, 0]), array([1, 0, 2, 3]))
# note: unraveling indices obtained by partitioning an array of shape (2, n) will not keep into account the first dimension (here [0,0,0,0])

# Craft batch indices...
batch_indices = np.repeat(np.arange(arr.shape[0], k)
# ...and join
maxk_indices = tuple([batch_indices]+[ind for ind in maxk_ind_shape[1:]])

# The result is used to re-assign k-highest values for each batch element to a destination matrix:
arr2 = np.zeros_like(arr)
arr2[maxk_indices] = arr[maxk_indices]
# arr2 prints: 
# array([[[ 0, 0, 0, 0],
#         [ 0, 0, 0, 0],
#         [23,22, 0, 0]],
#
#        [[ 0, 0, 14, 15],
#         [ 0, 0, 0, 0],
#         [ 0, 0, 0, 0]]])

任何帮助，将不胜感激。

标签： numpy

Answer 1

一种方法是使用np.[put/take]_along_axis：

gr_ind = np.argpartition(arr_batched,-k,axis=-1)[:,-k:]
arr_2 = np.zeros_like(arr)
np.put_along_axis(arr_2.reshape(arr_batched.shape),gr_ind,np.take_along_axis(arr_batched,gr_ind,-1),-1)