Numpy широковещательный массив для меньшего массива с точной позицией для каждой строки

Question

Рассмотрим пример матричного массива:

[[0 1 2 1 0]
 [1 1 2 1 0]
 [0 1 0 0 0]
 [1 2 1 0 0]
 [1 2 2 3 2]]

Что мне нужно сделать:

1. найти максимумы в каждом ряду
2. выбрать меньшее окружение максимумов из каждой строки (в данном случае 3 значения)
3. вставить окружение максимумов в новый массив (уже)

Для приведенного выше примера результат:

[[ 1.  2.  1.]
 [ 1.  2.  1.]
 [ 0.  1.  0.]
 [ 1.  2.  1.]
 [ 2.  3.  2.]]

Мой текущий рабочий код:

import numpy as np

A = np.array([
    [0, 1, 2, 1, 0],
    [1, 1, 2, 1, 0],
    [0, 1, 0, 0, 0],
    [1, 2, 1, 0, 0],
    [1, 2, 2, 3, 2],
])

b = A.argmax(axis=1)

C = np.zeros((len(A), 3))
for idx, loc, row in zip(range(len(A)), b, A):
    print(idx, loc, row)
    C[idx] = row[loc-1:loc+2]

print(C)

Мой вопрос:

Как избавиться от цикла for и заменить его на более дешевую операцию с клочками?

Примечание:

Этот алгоритм предназначен для выпрямления ломаных "линий" в кадрах видеопотока с тысячами строк.

Answer 1

Это работает путем создания и массива со всеми желаемыми индексами, но каким-то образом использование этого непосредственно на A приводит к созданию трехмерного массива, следовательно, последующего индексирования ... Вероятно, не оптимальное, но определенно другой способ сделать это!

import numpy as np

A = np.array([
    [0, 1, 2, 1, 0],
    [1, 1, 2, 1, 0],
    [0, 1, 0, 0, 0],
    [1, 2, 1, 0, 0],
    [1, 2, 2, 3, 2],
])

b = A.argmax(axis = 1).reshape(-1, 1)
index = b + np.arange(-1,2,1).reshape(1, -1)
A[:,index][np.arange(b.size),np.arange(b.size)]

Answer 2

Подход № 1

У нас может быть векторизованное решение, основанное на настройке скользящих окон и последующей индексации на них с индексами со смещением по b, чтобы получить желаемый результат. Мы можем использовать np.lib.stride_tricks.as_strided на основе scikit-image's view_as_windows, чтобы получить раздвижные окна. Подробнее об использовании as_strided на основе view_as_windows.

Реализация будет -

from skimage.util.shape import view_as_windows

L = 3 # window length
w = view_as_windows(A,(1,L))[...,0,:]
Cout = w[np.arange(len(b)),b-L//2]

Будучи методом, основанным на представлении, он имеет преимущество, заключающееся в том, что он экономит память и, следовательно, также хорош для производительности.

Подход № 2

В качестве альтернативы, однострочник путем создания всех этих индексов с внешним добавлением будет -

A[np.arange(len(b))[:,None],b[:,None] + np.arange(-(L//2),L//2+1)]