Select certain rows (condition met), but only some columns in Python/Numpy

numpy, python

Solution

>>> a = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12]])
>>> a
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])

>>> a[a[:,0] > 3] # select rows where first column is greater than 3
array([[ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])

>>> a[a[:,0] > 3][:,np.array([True, True, False, True])] # select columns
array([[ 5,  6,  8],
       [ 9, 10, 12]])

# fancier equivalent of the previous
>>> a[np.ix_(a[:,0] > 3, np.array([True, True, False, True]))]
array([[ 5,  6,  8],
       [ 9, 10, 12]])

For an explanation of the obscure `np.ix_()`, see https://stackoverflow.com/a/13599843/4323

Finally, we can simplify by giving the list of column numbers instead of the tedious boolean mask:

>>> a[np.ix_(a[:,0] > 3, (0,1,3))]
array([[ 5,  6,  8],
       [ 9, 10, 12]])

Problem

I have an numpy array with 4 columns and want to select columns 1, 3 and 4, where the value of the second column meets a certain condition (i.e. a fixed value). I tried to first select only the rows, but with all 4 columns via: ``` I = A[A[:,1] == i] ``` which works. Then I further tried (similarly to matlab which I know very well): ``` I = A[A[:,1] == i, [0,2,3]] ``` which doesn't work. How to do it? EXAMPLE DATA: ``` >>> A = np.array([[1,2,3,4],[6,1,3,4],[3,2,5,6]]) >>> print A [[1 2 3 4] [6 1 3 4] [3 2 5 6]] >>> i = 2 # I want to get the columns 1, 3 and 4 # for every row which has the value i in the second column. # In this case, this would be row 1 and 3 with columns 1, 3 and 4: [[1 3 4] [3 5 6]] ``` I am now currently using this: ``` I = A[A[:,1] == i] I = I[:, [0,2,3]] ``` But I thought that there had to be a nicer way of doing it... (I am used to MATLAB)

Original source

Related problems