Home > database >  Remove rows in a 2d-numpy array if they contain a specific element
Remove rows in a 2d-numpy array if they contain a specific element

Time:05-27

I have a matrix as 2d-np.array and I would like to remove all rows that contain an element x in a specific column. My goal is to return a matrix without these rows, so it should be smaller. My function looks like this:

def delete_rows(matrix, x, col):
    for i in range(matrix.shape[0]-1):
        if(matrix[i,col] == x):
            np.delete(matrix, i, axis = 0)
    return matrix

Sadly in my test the shape of the matrix stayed the same after removing rows. I think the deleted rows were substituted by rows with 0s. Any advice on how I can achieve my goal?

CodePudding user response:

This can be simply done in one line:

import numpy as np

def delete_rows(matrix, x, col):
    return matrix[matrix[:,col]!=x,:]

For example, if we want to remove all the rows that contain 5 in the second column from matrix A:

>>> A = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> print(delete_rows(A, 5, 1))
[[1 2 3]
 [7 8 9]]

CodePudding user response:

Assuming you have an array like this:

array([[12,  5,  0,  3, 11,  3,  7,  9,  3,  5],
       [ 2,  4,  7,  6,  8,  8, 12, 10,  1,  6],
       [ 7,  7, 14,  8,  1,  5,  9, 13,  8,  9],
       [ 4,  3,  0,  3,  5, 14,  0,  2,  3,  8],
       [ 1,  3, 13,  3,  3, 14,  7,  0,  1,  9],
       [ 9,  0, 10,  4,  7,  3, 14, 11,  2,  7],
       [12,  2,  0,  0,  4,  5,  5,  6,  8,  4],
       [ 1,  4,  9, 10, 10,  8,  1,  1,  7,  9],
       [ 9,  3,  6,  7, 11, 14,  2, 11,  0, 14],
       [ 3,  5, 12,  9, 10,  4, 11,  4,  6,  4]])

You can remove all rows containing a 3 like this:

row_mask = np.apply_along_axis(np.any, 0, arr == 3)
arr = arr[~row_mask]

Your new array looks like this

array([[ 7,  7, 14,  8,  1,  5,  9, 13,  8,  9],
       [12,  2,  0,  0,  4,  5,  5,  6,  8,  4],
       [ 1,  4,  9, 10, 10,  8,  1,  1,  7,  9],
       [ 3,  5, 12,  9, 10,  4, 11,  4,  6,  4]])
  • Related