CODEWITHSUNDEEP
pythonarraysnumpymultidimensional-arrayindexing
mathguy
mathguy

Reputation: 1518

argmax for multidimensional array along some axis

I have a multidimension array that looks like this:

my_array = np.arange(2)[:,None,None] *np.arange(4)[:, None]*np.arange(8)

I am looking for a multidimensional equivalent of the 2-D argmax

In particular, I am looking for argmax of maxima along axis = 2. I tried reshaping first, but reshaping will completely destroy the original indices information of the entire array, so it probably won't work. I have no clue how to do it and need helps from you guys. Thank you in advance

EDIT: Desire output is:

[(0,0,0),(1,3,1),(1,3,2),(1,3,3),(1,3,4),(1,3,5),(1,3,6),(1,3,7)]

This exactly is the array of the indices of maxima along axis = 2

Upvotes: 1

Views: 67

Answers (1)

Divakar
Divakar

Reputation: 221504

For finding such argmax indices along the last axis of a 3D ndarray, we can use something along these lines -

In [66]: idx = my_array.reshape(-1,my_array.shape[-1]).argmax(0)

In [67]: r,c = np.unravel_index(idx,my_array.shape[:-1])

In [68]: l = np.arange(len(idx))

In [69]: np.c_[r,c,l]
Out[69]: 
array([[0, 0, 0],
       [1, 3, 1],
       [1, 3, 2],
       [1, 3, 3],
       [1, 3, 4],
       [1, 3, 5],
       [1, 3, 6],
       [1, 3, 7]])

To extend this to a generic ndarray -

In [99]: R = np.unravel_index(idx,my_array.shape[:-1])

In [104]: np.hstack((np.c_[R],l[:,None]))
Out[104]: 
array([[0, 0, 0],
       [1, 3, 1],
       [1, 3, 2],
       [1, 3, 3],
       [1, 3, 4],
       [1, 3, 5],
       [1, 3, 6],
       [1, 3, 7]])

Upvotes: 1

Related Questions