Home > Software engineering >  Checking non zero-sum rows in numpy array and removing them
Checking non zero-sum rows in numpy array and removing them

Time:03-15

I have a numpy array like this:

array([[ 3.,  2.,  3., ...,  0.,  0.,  0.],
       [ 3.,  2., -4., ...,  0.,  0.,  0.],
       [ 3., -4.,  1., ...,  0.,  0.,  0.],
       ...,
       [-1., -2.,  4., ...,  0.,  0.,  0.],
       [ 4., -2., -2., ...,  0.,  0.,  0.],
       [-2.,  2.,  4., ...,  0.,  0.,  0.]], dtype=float32)

what I want to do is removing all the rows that do not sum to zero and remove them, while also saving such rows indexes/positions in order to eliminate them to another array.

I'm trying the following:

for i in range(len(arr1)):
    
    count=0
    
    for j in arr1[i]:
        count =j
        
    if count != 0:
        
        arr_1 = np.delete(arr1,i,axis=0)
        arr_2 = np.delete(arr2,i,axis=0)

the resulting arr_1 and arr_2 still contain rows that do not sum to zero. What am I doing wrong?

CodePudding user response:

You can compute sum then keep row that have sum == 0 like below:

a=np.array([
    [ 3.,  2.,  3.,  0.,  0.,  0.],
    [ 3.,  2., -4.,  0.,  0.,  0.],
    [ 3., -4.,  1.,  0.,  0.,  0.]])

b = a.sum(axis=1) 
# array([8., 1., 0.])

print(a[b==0])

Output:

array([[ 3., -4.,  1.,  0.,  0.,  0.]])

CodePudding user response:

Just use sum(axis=1):

mask = a.sum(axis=1) != 0
do_sum_to_0 = a[~mask]
dont_sum_to_0 = a[mask]
  • Related