How to convert a 2d numpy array to a binary indicator matrix for maximum value

Question

Assuming I have a 2d numpy array indicating the probabilities for m samples in n classes (probabilities add up to 1 for each sample).

Assuming that each sample can only be in one category, I want to create a new array with the same shape as the original, but only with binary values ​​indicating which class had the highest probability.

Example:

[[0.2, 0.3, 0.5], [0.7, 0.1, 0.1]] 

should be converted to:

 [[0, 0, 1], [1, 0, 0]] 

It seems that amax is already doing almost what I want, but instead of indexes, I need a matrix of indicators, as described above.

It seems simple, but somehow I can't figure it out using the standard numpy functions. Of course, I could use regular python loops, but it seems like there should be an easier way.

If multiple classes have the same probability, I would prefer a solution that selects only one of the classes (I don't care, in this case).

Thanks!

Answer 1

Here is one way:

 In [112]: a Out[112]: array([[ 0.2, 0.3, 0.5], [ 0.7, 0.1, 0.1]]) In [113]: a == a.max(axis=1, keepdims=True) Out[113]: array([[False, False, True], [ True, False, False]], dtype=bool) In [114]: (a == a.max(axis=1, keepdims=True)).astype(int) Out[114]: array([[0, 0, 1], [1, 0, 0]]) 

(But this will give a True value for each occurrence of the maximum in the row. See Divakar's answer for a good way to select only the first occurrence of the maximum.)

Answer 2

In the case of links (two or more elements are the highest in the row) where you want to select only one, here is one approach to do this with np.argmax and broadcasting -

 (A.argmax(1)[:,None] == np.arange(A.shape[1])).astype(int) 

Run Example -

 In [296]: A Out[296]: array([[ 0.2, 0.3, 0.5], [ 0.5, 0.5, 0. ]]) In [297]: (A.argmax(1)[:,None] == np.arange(A.shape[1])).astype(int) Out[297]: array([[0, 0, 1], [1, 0, 0]])