ocean800 ocean800 - 6 days ago 6
Python Question

Python - Select rows of array on certain condition?

So I have an 2D array

data
that looks like this:

enter image description here

I want to count the number of rows on a certain condition of the second two columns. For example, in this particular slice of the array, I only have

1 | 2


But given the third row is
range(1,4)
and the fourth
range(0,3)
, I could have all of the following combinations:

1 | 0
1 | 1
1 | 2

2 | 0
2 | 1
2 | 2

3 | 0
3 | 1
3 | 2


I want to select the rows for which each of those conditions is true. But I'm not sure how to go about it? I've been working on this for the last 2 hours and I've come up with things using for loops, list comprehensions, etc. But it just gets more and more complicated and none of those ways actually worked. Is there a good way to do this in numpy, or even just plain python?

Any help would be greatly appreciated, thanks!!

Answer

This works:

import numpy as np
# data array 
data = np.array([[4,3,1,2],[4,3,5,1],[1,2,1,0]])
# array of acceptable combinations
cond = np.array([[1,0],[1,2]])
# index of rows matching the conditions
idx=np.array([row in cond for row in data[:,2:]])
# selected rows
data[idx]
# array([[4, 3, 1, 2],
#   [1, 2, 1, 0]]