Как сделать двумерный условный поиск?

Question

Я пытаюсь векторизовать или иным образом ускорить (вероятно, используя numpy) цикл поиска / сопоставления для.Я изучил np.vectorize, индексирование NumPy и np.where, но не могу найти правильную реализацию / комбинацию, соответствующую моим потребностям.

Код вопроса:

Sx = np.zeros((Np+1, 2*N+1))
rows, cols = prepped_array.shape[0], prepped_array.shape[1]

for ind1 in range(rows):
    for ind2 in range(cols):
        if prepped_array[ind1][ind2][0] != -1:
            Sx[ind1, ind2] = M[prepped_array[ind1][ind2][0], prepped_array[ind1][ind2][1]]

prepped_array является справочной таблицей (инициализированной для всех [-1, -1]), где значения были заменены, где они должны быть изменены в Sx.

M - это преобразованный вход, который мы хотим отобразить в массив Sx.

Есть идеи / указатели?Спасибо!

Answer 1

Вы можете использовать логическую маску для индексации Sx и prepped_array, а затем использовать два индексных массива, полученных из prepped_array, для индексации в массив M. Код может быть понятнее предыдущего предложения:

mask = prepped_array[:, :, 0] != -1
Sx[mask] = M[tuple(prepped_array[mask].T)]

Давайте посмотрим на вовлеченные шаги:

• mask = prepped_array[:, :, 0] != -1 создает двумерный логический массив, указывающий, где выполняется условие.
• prepped_array[mask] создает двумерный массив, в котором записи из предыдущих 3-х измерений теперь отображаются вдоль 2-го измерения; первые измерения соответствуют каждому True экземпляру в mask.
• tuple(prepped_array[mask].T) создает два одномерных массива, которые можно использовать для дальнейшего индексирования в другие массивы: первый массив обозначает индексы строк, а второй массив обозначает индексы столбцов.
• Итак, Sx[mask] = M[tuple(prepped_array[mask].T)] отображает индексы, содержащиеся в prepped_array, в массив M, используя два предыдущих массива индексов 1D.
• Sx[mask], наконец, ссылается на те элементы в Sx, для которых выполняется условие в prepped_array[:, :, 0].