Filtering pandas dataframe using element of list

user_12998544 0
Viewed 79

I am trying to filter a dataframe where some columns are lists. And I want to base the filter out elements that does not pass the condition.

For example:

import pandas as pd 
df = pd.DataFrame({'col1':[10,20], 'col2': [[1,2,3],[3,4,5]], 'col3': [[False,False,True],[True,True,False]],'col4':[True,False]})

   col1       col2                  col3   col4
0    10  [1, 2, 3]  [False, False, True]   True
1    20  [3, 4, 5]   [True, True, False]  False

applying the filter

df_filtered = df.query("col2>2 & col3==True")

the output I expect

enter image description here

Thanks for the help!

pandas python
4 Answers

Try:

df[['col2','col3']] = (pd.DataFrame({'col2': df['col2'].explode(),
                                     'col3': df['col3'].explode()})
                         .query('col2>2 & col3==True')
                         .groupby(level=0).agg(list)
                      )

Output:

print(df)

   col1    col2          col3   col4
0    10     [3]        [True]   True
1    20  [3, 4]  [True, True]  False
user_42384080

You can use numpy and an iterative approach if memory is the main constraint.

This modifies the dataframe in place without having to create a large interim data structure in the process:

import pandas as pd
import numpy as np

df = pd.DataFrame({'col1':[10,20], 'col2': [[1,2,3],[3,4,5]], 'col3': [[False,False,True],[True,True,False]]})

for idx, row in df.iterrows():
    a1=np.array(row['col2'])
    a2=np.array(row['col3'])
    df.at[idx,'col2']=a1[(a1>2) & a2]
    df.at[idx,'col3']=a2[a2]

>>> df
   col1    col2          col3
0    10     [3]        [True]
1    20  [3, 4]  [True, True]
user_2986070

As lists are same size across the rows, you can probably use arrays and mask like this

arr2 = np.array(df['col2'].tolist())
arr3 = np.array(df['col3'].tolist())

df[['col2','col3']] = [[c2[b],c3[b]] for c2,c3,b in zip(arr2,arr3,(arr2>=2) & arr3)]

print(df)
   col1    col2          col3   col4
0    10     [3]        [True]   True
1    20  [3, 4]  [True, True]  False
user_92747320

Another way with loops, but probably slower:

for index, row in df.iterrows():
    j=0
    for i in df.at[index, 'col3']:
        if i==False:
            df.at[index, 'col2'].remove(df.at[index, 'col2'][j])
        else:
            j=j+1
    df.at[index, 'col3']=list(filter(None, df.at[index, 'col3']))
user_97142550
Related