Как эффективно отфильтровать максимум элементов матрицы в строке

Question

Учитывая двумерный массив, я ищу питонный способ получить массив одинаковой формы, с максимальным элементом для каждой строки. См. Функцию max_row_filter ниже

def max_row_filter(mat2d):
    m = np.zeros(mat2d.shape)
    for r in range(mat2d.shape[0]):
        c = np.argmax(mat2d[r])
        m[r,c]=mat2d[r,c]
    return m

p = np.array([[1,2,3],[5,4,3,],[9,10,3]])
max_row_filter(p)

Out: array([[ 0.,  0.,  3.],
            [ 5.,  0.,  0.],
            [ 0., 10.,  0.]])

Я ищу эффективный способ сделать это, подходящий для больших массивов.

Answer 1

Альтернативный ответ (сохранятся дубликаты):

p * (p==p.max(axis=1, keepdims=True))

Answer 2

Если дубликатов нет, вы можете использовать numpy.argmax :

import numpy as np

p = np.array([[1, 2, 3],
              [5, 4, 3, ],
              [9, 10, 3]])

result = np.zeros_like(p)

rows, cols = zip(*enumerate(np.argmax(p, axis=1)))
result[rows, cols] = p[rows, cols]

print(result)

выход

[[ 0  0  3]
 [ 5  0  0]
 [ 0 10  0]]

Обратите внимание, что для нескольких случаев argmax возвращает первый случай.