Selecting rows from a NumPy ndarray

Question

I want to select only certain rows from a NumPy array based on the value in the second column. For example, this test array has integers from 1 to 10 in the second column.

Answer 1

test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]

The result should be faster than the original version since NumPy's doing the inner loops instead of Python.

Answer 2

numpy.in1d is what you are looking for:

print test[numpy.in1d(test[:,1], wanted)]

It should easily be the fastest solution if wanted is large; plus, it is the most readable one, id say.

Answer 3

This is two times faster than Amnon's variant for len(test)=1000:

wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]

Answer 4

The following solution should be faster than Amnon's solution as wanted gets larger:

# Much faster look up than with lists, for larger lists:
wanted_set = set(wanted)

@numpy.vectorize
def selected(elmt): return elmt in wanted_set
# Or: selected = numpy.vectorize(wanted_set.__contains__)

print test[selected(test[:, 1])]

In fact, it has the advantage of searching through the test array only once (instead of as many as len(wanted) times as in Amnon's answer). It also uses Python's built-in fast element look up in sets, which are much faster for this than lists. It is also fast because it uses Numpy's fast loops. You also get the optimization of the in operator: once a wanted element matches, the remaining elements do not have to be tested (as opposed to the "logical or" approach of Amnon, were all the elements in wanted are tested no matter what).

Alternatively, you could use the following one-liner, which also goes through your array only once:

test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]

This is much much slower, though, as this extracts the element in the second column at each iteration (instead of doing it in one pass, as in the first solution of this answer).