1

Suppose I have the following numpy array.

Q = np.array([[0,1,1],[1,0,1],[0,2,0])

Question: How do I identify the position of each max value along axis 1? So the desired output would be something like:

array([[1,2],[0,2],[1]]) # The dtype of the output is not required to be a np array.

With np.argmax I can identify the first occurrence of the maximum along the axis, but not the subsequent values.

In: np.argmax(Q, axis =1) 
Out: array([1, 0, 1])

I've also seen answers that rely on using np.argwhere that use a term like this.

np.argwhere(Q == np.amax(Q))

This will also not work here because I can't limit argwhere to work along a single axis. I also can't just flatten out the np array to a single axis because the max's in each row will differ. I need to identify each instance of the max of each row.

Is there a pythonic way to achieve this without looping through each row of the entire array, or is there a function analogous to np.argwhere that accepts an axis argument?

Any insight would be appreciated thanks!

phntm
  • 511
  • 2
  • 11

1 Answers1

1

Try with np.where:

np.where(Q == Q.max(axis=1)[:,None])

Output:

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

Not quite the output you want, but contains equivalent information.

You can also use np.argwhere which gives you the zip data:

np.argwhere(Q==Q.max(axis=1)[:,None])

Output:

array([[0, 1],
       [0, 2],
       [1, 0],
       [1, 2],
       [2, 1]])
Quang Hoang
  • 146,074
  • 10
  • 56
  • 74
  • oh awesome thanks alot! This is a really basic question but would you mind explaining what function the [:, None] is playing? Still trying to get the hang of this. Also I just want to clarify is there a distinction between max and amax? – phntm Mar 04 '21 at 20:17
  • 1
    @phntm `[:,None]` thing is call [broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html#:~:text=The%20term%20broadcasting%20describes%20how%20numpy%20treats%20arrays,that%20looping%20occurs%20in%20C%20instead%20of%20Python.). `.max` and `.amax` are [the same thing](https://stackoverflow.com/questions/33569668/numpy-max-vs-amax-vs-maximum). – Quang Hoang Mar 04 '21 at 20:34