LetsPlayYahtzee LetsPlayYahtzee - 7 months ago 24
Python Question

Apply a comparison between two numpy.arrays to only one column but retrieve whole rows

I have two numpy arrays with two columns each.

a = numpy.array([[1131, 1],
[4131, 2],
[421, 1],
[41, 1]])

b = numpy.array([[5881, 2],
[637, 2],
[742, 2],
[36, 2]])


and I want to create a third array with 2 columns that the fist column will contain the minimum between the first columns of a and b and the second column will contain whatever the second column of the array containing the minimum had.

the third array should be

c = numpy.array([[1131, 1],
[637, 2],
[421, 1],
[36, 2]])


How can I do that efficiently?

Answer

Looks like a perfect case to ab-use NumPy broadcasting within np.where -

np.where((a[:,0] < b[:,0])[:,None],a,b)

The beauty is that it would work independent of the number of columns in a and b, as that's where the broadcasting helps out.

Sample run -

In [78]: a
Out[78]: 
array([[1131,   99],
       [4131,    4],
       [ 421,   56],
       [  41,   78]])

In [79]: b
Out[79]: 
array([[5881,   23],
       [ 637,   42],
       [ 742, 7882],
       [  36,   62]])

In [80]: np.where((a[:,0] < b[:,0])[:,None],a,b)
Out[80]: 
array([[1131,   99],
       [ 637,   42],
       [ 421,   56],
       [  36,   62]])