numpy: Чистый выбор координат (индексов) для самых высоких значений k - вдоль указанной оси c - в 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]]])

Любая помощь будет признательна.

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)