Перевод функции max в MatLab в numpy

Question

Я хочу перевести следующий код MatLab в Numpy.

% MatLab
[selVals, colAssignment] = max(A, [], 2);

Согласно руководству пользователя NumPy для Matlab на сайте Numpy, синтаксис для преобразования max (..., [], 2) =

# Numpy
selVals = A.max(1)

Однако, это просто дает простой эквивалент selVals. Это не дает назначение. Как мне получить colAssignment в Numpy?

Answer 1

Предполагая, что A - это массив numpy, вы можете сделать это в одной строке кода.

selVals, colAssignment = A.max(axis=1), A.argmax(axis=1)

В качестве альтернативы:

colAssignment = A.argmax(axis=1) #index
selVals = A[:, colAssignment] #value

Answer 2

Вы не можете сделать это за один вызов функции. Чтобы получить colAssignment, позвоните argmax:

setVals = A.max(axis=1)
colAssignment = A.argmax(axis=1)

В зависимости от размера A может быть более эффективным сначала получить colAssignment и использовать его для извлечения максимальных значенийот A:

colAssignment = A.argmax(axis=1)
setVals = A[np.arange(len(A)), colAssignment]