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

前端 未结 5 2068
一整个雨季
一整个雨季 2020-12-15 05:29

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 f

5条回答
  •  萌比男神i
    2020-12-15 06:05

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

提交回复
热议问题