Selecting rows from 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 second column.

>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose()
>>> test[:10, :]
array([[ 0,  6],
       [ 1,  7],
       [ 2, 10],
       [ 3,  4],
       [ 4,  1],
       [ 5, 10],
       [ 6,  6],
       [ 7,  4],
       [ 8,  6],
       [ 9,  7]])

If I wanted only rows where the second value is 4, it is easy.

>>> test[test[:, 1] == 4]
array([[ 3,  4],
       [ 7,  4],
       [16,  4],
       ...
       [81,  4],
       [83,  4],
       [88,  4]])

But how do I achieve the same result when there is more than one wanted value. The wanted list can be of arbitrary length. For example, I may want all rows where the second column is either 2, 4 or 6.

>>> wanted = [2, 4, 6]

The only way I have come up with is to use list comprehension and then convert this back into an array and seems too convoluted, although it works.

>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])]
array([[ 0,  6],
       [ 3,  4],
       [ 6,  6],
       ...
       [90,  2],
       [91,  6],
       [92,  2]])

Is there a better way to do this in numpy itself that I am missing?

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

The following solution has the advantage of going through your array only once:

@numpy.vectorize
def selected(elmt): return elmt in wanted

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

It is also fast because it uses Numpy's fast loops. You also get the optimization of the in operator: once an element matches, the remaining elements do not have to be tested (as opposed to the "logical or" approach, were all the elements in wanted are tested, possibly unnecessarily).

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 is in one pass, as in the first solution).

Answer 3

This is ten 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), :]