Определить индекс наивысшей ценности в Python NumPy

Я хочу создать массив с индексом максимального максимального значения каждой строки.

a = np.array([ [1,2,3], [6,5,4], [0,1,0] ])
maxIndexArray = getMaxIndexOnEachRow(a)
print maxIndexArray 

[[2], [0], [1]]

Там есть функция np.argmax, но она, похоже, не делает то, что я хочу...

Ответ 1

Функция argmax() выполняет то, что вы хотите:

print a.argmax(axis=1)
array([2, 0, 1])