find argmin of 3 vectors in numpy-CodePudding

I have three vectors
a = np.arange(3)
b = np.arange(2,-1,-1)
c = np.ones((3,))
and their minimum values element-wise via
np.minimum(a, b, c)
which is: m = [0. 1. 0.]
I want to find from which array these values came, in my example I want input like: [0. x. 1.] because first element in m came from a and last element in m came from b

CodePudding user response：

You can use the argmin method with the axis argument:

>>> import numpy as np
>>> a = np.arange(3)
>>> b = np.arange(2,-1,-1)
>>> c = np.ones((3,))
>>> a
array([0, 1, 2])
>>> b
array([2, 1, 0])
>>> c
array([1., 1., 1.])
>>> all = np.stack([a, b, c])
>>> all
array([[0., 1., 2.],
       [2., 1., 0.],
       [1., 1., 1.]])
>>> all.min(axis=0)
array([0., 1., 0.])
>>> all.argmin(axis=0)
array([0, 0, 1])