How to get indexes of all maximum values in python multidimensional np array

user_14423481 0
Viewed 45

I would like to get all indexes of maximum values for each row from ndarray. For example, i have

arr = np.array([[1, 3, 3], [1, 5, 4]])

And i would like to get indexes of all 3's from first row, and all 5's from second row.

np.where(((arr == arr[0].max()) | (arr == arr[1].max())))

And it returns

(array([0, 0, 1], dtype=int64), array([1, 2, 1], dtype=int64))

I want something like that but more universal for any amount of rows. Because np.where(arr == arr.argmax()) doesn't work like i want it to work. It only returns indexes of first maxium that it met in each row.

numpy arrays python
1 Answers

@Paul's answer in comments might be the best you can find. Writing it for readers. arr.max(1) find max in each row and arr==arr.max(1,keepdims=True) finds all elements in each row that are equal to corresponding max in that row. Finally nonzero returns indices of those elements:

(arr==arr.max(axis=1,keepdims=True)).nonzero()

output for OP's example:

(array([0, 0, 1]), array([1, 2, 1]))
user_49759810
Related