## Return the subset of NumPy array according to the first element of each row

Question

I am trying to get the subset *x* of the given NumPy array *alist* such that the first element of each row must be in the list *r*.

```
>>> import numpy
>>> alist = numpy.array([(0, 2), (0, 4), (1, 3), (1, 4), (2, 1), (3, 1), (3, 2), (4, 1), (4, 3), (4, 2)])
>>> alist
array([[0, 2],
[0, 4],
[1, 3],
[1, 4],
[2, 1],
[3, 1],
[3, 2],
[4, 1],
[4, 3],
[4, 2]])
>>> r = [1,3]
>>> x = alist[where first element of each row is in r] #this i need to figure out.
>>> x
array([[1, 3],
[1, 4],
[3, 1],
[3, 2]])
```

Any easy way (without looping as I've a large dataset) to do this in Python?

Show source

## Answers to Return the subset of NumPy array according to the first element of each row ( 2 )

Slice the first column off input array (basically selecting first elem from each row), then use

`np.in1d`

with`r`

as the second input to create a mask of such valid rows and finally index into the rows of the array with the mask to select the valid ones.Thus, the implementation would be like so -

Sample run -

You can construct the index array for the valid rows using some indexing tricks: we can add an additional dimension and check equality with each element of your first column:

The trick is that we take the first column of

`alist`

, make it an`(N,1)`

-shaped array, make use of array broadcasting in the comparison to end up with an`(N,2)`

-shape boolean array, and if any of the values in a given row is`True`

, we keep that index. The resulting index array is the exact same as the`np.in1d`

one in Divakar's answer.