Element wise comparison between 1D and 2D array

Question

Want to perform an element wise comparison between an 1D and 2D array. Each element of the 1D array need to be compared (e.g. greater) against the corresponding row of 2D and a mask will be created. Here is an example:

A = np.random.choice(np.arange(0, 10), (4,100)).astype(np.float)
B = np.array([5., 4.,  8.,  2. ])

I want to do

A<B 

so that first row of A will be compared against B[0] which is 5. and the result will be an boolean array.

If I try this I get:

operands could not be broadcast together with shapes (4,100) (4,)

Any ideas?

Answer 1

You need to insert an extra dimension into array B:

A < B[:, None]

This allows NumPy to properly match up the two shapes for broadcasting; B now has shape (4, 1) and the dimensions can be paired up:

(4, 100)
(4,   1)

The rule is that either the dimensions have the same length, or one of the lengths needs to be 1; here 100 can be paired with 1, and 4 can be paired with 4. Before the new dimension was inserted, NumPy tried to pair 100 with 4 which raised the error.