CODEWITHSUNDEEP
pythonarrayslistnumpy
Link_tester
Link_tester

Reputation: 1081

how to filter arrays stored as two lists

I have to lists of arrays, one is all my data and another one is a filter:

data= [array([[12., 2., 7.],
              [9., -1., 0.],
              [2., 0., 4.],
              [14., 2., 1.2]]),
       array([[11., 3., 1.],
              [2., 5., 0.],
              [1., -1., 1.],
              [-8., 0., 6.],
              [2., 3., 3.]])]

filter= [array([2., 0., 5.]),
         array([3., -1.])]

I want to only keep values of each array of my data, that their second column (data[i][:,1]) is equal to one the values of the same array from the filter (all values of filter[i]). I want to have the following output:

out = [array([[12., 2., 7.],
              [2., 0., 4.],
              [14., 2., 1.2]]),
       array([[11., 3., 1.],
              [1., -1., 1.],
              [2., 3., 3.]])]

I have tried the following code but it was not successful:

out=[]
for a in data:
    for b in filter:
        for c in b:
            if a[:,1]==c:
                bb=a[:,c]
                out=np.append(data,bb)

I appreciate any help and contribution in advance.

Upvotes: 0

Views: 880

Answers (1)

Dani Mesejo
Dani Mesejo

Reputation: 61920

You could use np.isin to create the masks and then filter using those masks:

import numpy as np

data = [np.array([[12., 2., 7.],
                  [9., -1., 0.],
                  [2., 0., 4.],
                  [14., 2., 1.2]]),
        np.array([[11., 3., 1.],
                  [2., 5., 0.],
                  [1., -1., 1.],
                  [-8., 0., 6.],
                  [2., 3., 3.]])]

filter_data = [np.array([2., 0., 5.]),
               np.array([3., -1.])]

# compute the masks
masks = [np.isin(arr[:, 1], filter_arr) for arr, filter_arr in zip(data, filter_data)]

# filter using the masks
result = [arr[mask] for arr, mask in zip(data, masks)]

for filtered_arr in result:
    print(filtered_arr)

Output

[[12.   2.   7. ]
 [ 2.   0.   4. ]
 [14.   2.   1.2]]
[[11.  3.  1.]
 [ 1. -1.  1.]
 [ 2.  3.  3.]]

Upvotes: 1

Related Questions