CODEWITHSUNDEEP
pythonnumpy
Mark
Mark

Reputation: 1069

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.

Upvotes: 1

Views: 68

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114440

You can extend the approach in the other answer to get the values of a and b. Given a mask of

c = b > a[:, None]

You can extract the indices for each dimension using np.where or np.nonzero:

row, col = np.nonzero(c)

And use the indices to get the corresponding values:

ag = a[row]
bg = b[col]

Elements of a and b may be repeated in the result.

Upvotes: 0

cottontail
cottontail

Reputation: 23271

If you're asking how to get c without loop, try this

# make "a" a column vector
# > broadcasts to produce a len(a) x len(b) array
c = b > a[:, None]
c
array([[False,  True,  True,  True],
       [False, False,  True,  True],
       [False, False, False,  True]])

Upvotes: 3

Related Questions