CODEWITHSUNDEEP
pythonnumpy
Bill Zack
Bill Zack

Reputation: 21

numpy finding indices of good values

I am a newbie in Numpy. I want to find indices of elements equal to a set of values. For example, I know this works:

>>> x = np.array([[0, 1], [2, 3], [4, 5]])
>>> x
array([[0, 1],
       [2, 3],
       [4, 5]])
>>> x[np.where(x[:,0] == 2)]
array([[2, 3]])

But why doesn't this work? I would imagine it should be a straight-forward extension. No?

>>> x[np.where(x[:,0] == [2,4])]
array([], shape=(0, 2), dtype=int64)

Upvotes: 2

Views: 127

Answers (3)

pseudocubic
pseudocubic

Reputation: 1039

From your question, I assume that you meant that you want to find rows where the first element is equal to some special value. The way you have stated it I think that you meant x[:,0] in [2,4], not x[:,0] == [2,4], which will still not work in np.where(). Instead, you would need to construct it something like this:

x[np.where((x[:,0] == 2) | (x[:,0] == 4))]

Since this doesn't scale well, you might instead want to try doing it in a for loop:

good_row = ()
good_values = [2,4]
for val in good_values:
    good_row = np.append(good_row,np.where(x[:,0] == val))

print x[good_row]

Upvotes: 0

user2357112
user2357112

Reputation: 281529

== doesn't work like that; when given two arrays of different shape, it tries to broadcast the comparison according to the broadcasting rules.

If you want to determine which elements of an array are in a list of elements, you want in1d:

>>> x = numpy.arange(9).reshape((3, 3))
>>> numpy.in1d(x.flat, [2, 3, 5, 7])
array([False, False,  True,  True, False,  True, False,  True, False], dtype=boo
l)
>>> numpy.in1d(x.flat, [2, 3, 5, 7]).reshape(x.shape)
array([[False, False,  True],
       [ True, False,  True],
       [False,  True, False]], dtype=bool)

Upvotes: 1

Joran Beasley
Joran Beasley

Reputation: 114038

you would want

 x[np.where(x[:,0] in [2,4])]

which may not actually be valid numpy

mask = numpy.in1d(x[:,0],[2,4])
x[mask]

as an aside you could rewrite your first condition to

x[x[:,0] == 2]

Upvotes: 0

Related Questions