Search code examples
pythonnumpyscipysoftmax

Tunring an array of probabilities, into a binary array

I am trying to convert a NumPy array of probabilities that sum up to 1.0, into a binary array. Essentially, I want to turn the max probability in the array into 1, and the rest into 0s. Of course, I know how to do this using Python, but I was hoping that NumPy or SciPy already come up with a built-in function for making the operation more performant (we are speaking about a matrix containing millions of such arrays).

Solution

  • Compare items with the maximum value, and cast to int

    >>> a=np.array([0.435,0.24,0.241,0.13,0.56])

>>> np.int8( a == a.max() )
array([0, 0, 0, 0, 1], dtype=int8)