comparing vector to each element of another vector with numpy-CodePudding

I have 2 vectors, I would like to compute a matrix where the i'th row is the first vector compared to the i'th element of the second vector. This code does that:

a1 = np.array([1, 2, 3])
a2 = np.array([1, 2, 3])

print(np.array([a1 > e for e in a2]))

where we get:

[[False  True  True]
 [False False  True]
 [False False False]]

I want to have the same behavior but done efficiently with "numpy magic". How can I do this?

CodePudding user response：

You can broadcast the function by making one of the arrays have one singleton dimension:

In [1]: import numpy as np                                                      

In [2]: a1 = np.array([1, 2, 3])
   ...: a2 = np.array([1, 2, 3])[:,None]                                        

In [3]: a1>a2                                                                   
Out[3]: 
array([[False,  True,  True],
       [False, False,  True],
       [False, False, False]])

CodePudding user response：

The np.meshgrid (see documentation) is useful for constructing 2-d arrays, where each of two 1d vectors is broadcast across the number of elements of the other. It might be easier to see what is going on if the arrays have different elements and sizes:

>>> a1 = np.array([1,5,7])
>>> a2 = np.array([2,3,4,5])

>>> a1_2d, a2_2d = np.meshgrid(a1, a2)

>>> a1_2d
array([[1, 5, 7],
       [1, 5, 7],
       [1, 5, 7],
       [1, 5, 7]])

>>> a2_2d
array([[2, 2, 2],
       [3, 3, 3],
       [4, 4, 4],
       [5, 5, 5]])

Note that these two vectors have been broadcast across different dimensions.

They can then be compared element-by-element using the > operator to give an array of booleans.

>>> a1_2d > a2_2d
array([[False,  True,  True],
       [False,  True,  True],
       [False,  True,  True],
       [False, False,  True]])

Or going back to your original vectors, this gives:

array([[False,  True,  True],
       [False, False,  True],
       [False, False, False]])