将MatLab中的max函数转换为numpy

Question

我想将下面的MatLab代码转换为Numpy。

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

根据Numpy网站上的《用于Matlab的NumPy用户指南》，用于转换max（...，[]，2）的语法是

# Numpy
selVals = A.max(1)

但是，这只是selVals的numpy等效项。它不提供colAssignment。我如何在Numpy中获得colAssignment？

Answer 1

您无法在单个函数调用中完成此操作。要获取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]