Numpy array comprehensions, where problem

Question

I have a little problem with np.where and I can't find any answers in google so here's my question. Consider I have these arrays:

a = np.array([[1, 0.5, 2], [5, 0.25, 9], [-1, 0.77, 3]])
b = np.array([0.6, 0.2, 0.6])

What I want is: if a[x, 1] < b[x] then a[x, 0] else a[x, 2] . So I want to get this:

[1, 9, 3]

I thought code like this would work, but it doesn't

np.where(a[1] < b, a[0], a[2])

Answer 1

You almost got it:

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

You need to specify which axis you want to make the comparisons and take the elements, else as a default it will use the 0 axis, check the values for:

a[1]
a[1,:]
a[:,1]

See that the first two are the same (the default behaviour), and you are looking for the third.