Skip to content
Advertisement

How to find maximums’ indexes by row with Numpy?

I have np 2d array and want to get indexes of max by row.

For example:

3 2 1 3 1
0 4 4 1 2
1 0 1 1 2

Maximum by row is [3, 4, 2]. Indexes are [(0,0) (0,3) (1,1) (1,2) (2,4)]

I tried smth like this

buf = np.apply_along_axis(lambda x: zip(np.where(x == np.max(x))), axis=1, arr=B)

or this

buf = np.where(B == np.max(B,axis = 1))

But it doesnt work.

Also I can’t use loops (otherwise I wouldn’t use numpy either)

Advertisement

Answer

After finding max you need set each max as one row then searching in base array like below:

>>> a = np.array([[3 ,2 ,1 ,3 ,1],[0, 4 ,4 ,1 ,2],[1, 0, 1, 1, 2]])

>>> np.max(a, axis=1)
array([3, 4, 2])

>>> np.max(a, axis=1)[:,None]
array([[3],
       [4],
       [2]])

>>> np.argwhere(a==np.max(a,1)[:,None])
array([[0, 0],
       [0, 3],
       [1, 1],
       [1, 2],
       [2, 4]])

If you want output like as tuple in your question you can try this:

>>> list(map(tuple, np.argwhere(a==np.max(a,1)[:,None])))
[(0, 0), (0, 3), (1, 1), (1, 2), (2, 4)]
User contributions licensed under: CC BY-SA
7 People found this is helpful
Advertisement