Search code examples
pythonlistnumpymin

Returning binary '1' at index of minimum value

I have an np array with shape (50,2)

so for each row, I'm trying to get the index of the minimum value which I think I managed to roughly

for i in range(len(mylist)):
list(my_list[i]).index(min(my_list[i]))

however, I am getting a brain block on how to insert '1' at the index of the minimum value?

for e.g

([[2,4],
 [5,3]])

will give the index values for the  min(my_list) as 
0
1

then the first row will have the index being 0 and the 2nd row will be 1

how do I insert a binary value at the index of the min value so that the output is like

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

thank you!

Solution

  • Use numpy.arange for the indices of the rows and numpy.argmin for the indices of the columns:

    import numpy as np

arr = np.array([[2, 4],
                [5, 3]])

res = np.zeros_like(arr)
res[np.arange(arr.shape[0]), arr.argmin(axis=1)] = 1
print(res)

    Output

    [[1 0]
 [0 1]]