How to convert 2d numpy array into binary indicator matrix for max value

Viewed 3107

assuming I have a 2d numpy array indicating probabilities for m samples in n classes (probabilities sum to 1 for each sample).

Assuming each sample can only be in one category, I want to create a new array with the same shape as the original, but with only 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 amax already does almost what I want, but instead of the indices I want an indicator matrix as descrived above.

Seems simple, but somehow I can't figure it out using standard numpy functions. I could use regular python loops of course, but it seems there should be a simpler way.

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

Thanks!

2 Answers
Related