How to get row numbers of maximum elements in a 2D Numpy array?

Question

I have a 2D array a given by:

a = np.array([[2, 3, 1, 9], [0, 5, 4, 7], [2, 4, 6, 8]])

[[2 3 1 9]
 [0 5 4 7]
 [2 4 6 8]]

I would like to get row numbers of maximum elements column-wise, i.e. given by np.amax(a, axis=0), which in my example are [0, 1, 2, 0]. What is the Numpy magic to accomplish it?

Answer 1

I believe you're looking for np.argmax.

print(np.argmax(a, axis=0))

Result:

[0 1 2 0]