25

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.

>>> 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?

Improve this question

4 Answers 4

Reset to default
32

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).

Improve this answer
Sign up to request clarification or add additional context in comments.

4 Comments

These solutions call python for every element instead of using numpy's comparison. According to my tests, your first solution is faster than mine for len(wanted)=50 but slower for len(wanted)=5.
EOL, many thanks for your time and effort. Your explanations were clear. I chose to use Amnon's solution because for my usual scenario (len(test) about 1000 and len(wanted) about 3-5), that was faster than your first solution. The speed difference is not huge, but I also found it clearer. But it was good to be reminded of numpy's vectorize and I am sure I will find a use for it soon.
What if I want to select certain rows with conditional values in column ? For eg, I want to select all the rows whose column 1 has values in the range of 5 to 10. Eg. test[test[:, 1] > 4 & test[:, 1] < 8].
You got it right, except that you need parentheses: test[(…) & (…)], because & has priority over the comparison operator <.
16 
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.

Improve this answer

3 Comments

This solution goes through the array len(wanted) times. It is usually faster to go through the array in a single pass.
Thanks Amnon. This is the solution that I decided to accept. I think it is clear to understand and is about 20 x faster than my original solution.
What if I want to select certain rows with conditional values in column ? For eg, I want to select all the rows whose column 1 has values in the range of 5 to 10. Eg. test[test[:, 1] > 4 & test[:, 1] < 8].
11

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.

Improve this answer

Comments

0

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), :]
Improve this answer

3 Comments

@ahatchkins: Some typos in your version. What you are suggesting is this - # convert wanted list into a one column array wanted = numpy.array(wanted).reshape((len(wanted),1)) # print test[numpy.any(test[:,1] == wanted, 0)] In my tests, this is about 2 times faster than Amnon's solution.
Yes, there was a typo: s/wanted2/wanted/ . Fixed
Hmm, yes. You are right. Two typos in fact. And only 2 times faster. )

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.