pythonnumpyconditional-statementsswitch-statementarray-broadcasting

Is there a 2-D "where" in numpy?

This might seem an odd question, but it boils down to quite a simple operation that I can't find a numpy equivalent for. I've looked at np.where as well as many other operations but can't find anything that does this:

a = np.array([1,2,3])
b = np.array([1,2,3,4])
c = np.array([i<b for i in a])

The output is a 2-D array (3,4), of booleans comparing each value.

Solution

  • If you're asking how to get c without loop, then you can make a a column vector and perform the comparison. Numpy will broadcast a and b to create a len(a)xlen(b) array of boolean values.

    c = b > a[:, None]

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