CODEWITHSUNDEEP
pythonpandasdataframegroup-bypandas-groupby
Zmnako Awrahman
Zmnako Awrahman

Reputation: 548

group by columns then filter based on a condition

I have a df and I want to filter out a column based on a grouping. I want to keep group by combinations ((cc, odd, tree1, and tree2) if day > 4, then keep it, otherwise drop it

df = pd.DataFrame()
df['cc'] = ['BB', 'BB', 'BB', 'BB','BB', 'BB','BB', 'BB', 'DD', 'DD', 'DD', 'DD', 'DD', 'DD', 'DD', 'DD', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ']
df['odd'] = [3434, 3434, 3434, 3434, 3435, 3435, 3435, 3435, 3434, 3434, 3434, 3434, 3435, 3435, 3435, 3435, 3434, 3434, 3434, 3434, 3435, 3435, 3435, 3435]
df['tree1'] = ['ASP', 'ASP', 'ASP', 'ASP', 'SAP', 'SAP', 'SAP', 'SAP', 'ASP', 'ASP', 'ASP', 'ASP', 'SAP', 'SAP', 'SAP', 'SAP', 'ASP', 'ASP', 'ASP', 'ASP', 'SAP', 'SAP', 'SAP', 'SAP']
df['tree2'] = ['ATK', 'ATK','ATK','ATK','ATK','ATK','ATK','ATK', 'ATK', 'ATK','ATK','ATK','ATK','ATK','ATK','ATK', 'ATK', 'ATK','ATK','ATK','ATK','ATK','ATK','ATK']
df['day'] = [1, 2, 3, 4, 3, 4, 5, 6, 2, 3, 4, 5, 1, 3, 5, 7, 1, 2, 6, 8, 2, 4, 6, 8]
df

I tried this but this drops any row with day value smaller than 4

df_grouped = df.groupby(['cc', 'odd', 'tree1', 'tree2']).filter(df['day'] > 4)

I get this error TypeError: 'Series' object is not callable

And tried this

df_grouped = df.groupby(['cc', 'odd', 'tree1', 'tree2']).filter(lambda x: x['day'] > 4)

I get this error TypeError: filter function returned a Series, but expected a scalar bool.

I searched and tried to solve these errors but the proposed solution did not work for me. I would like to get a df as below:

df1 = pd.DataFrame()
df1['cc'] = ['BB', 'BB','BB', 'BB', 'DD', 'DD', 'DD', 'DD', 'DD', 'DD', 'DD', 'DD', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ', 'ZZ']
df1['odd'] = [3435, 3435, 3435, 3435, 3434, 3434, 3434, 3434, 3435, 3435, 3435, 3435, 3434, 3434, 3434, 3434, 3435, 3435, 3435, 3435]
df1['tree1'] = ['SAP', 'SAP', 'SAP', 'SAP', 'ASP', 'ASP', 'ASP', 'ASP', 'SAP', 'SAP', 'SAP', 'SAP', 'ASP', 'ASP', 'ASP', 'ASP', 'SAP', 'SAP', 'SAP', 'SAP']
df1['tree2'] = ['ATK','ATK','ATK','ATK', 'ATK', 'ATK','ATK','ATK','ATK','ATK','ATK','ATK', 'ATK', 'ATK','ATK','ATK','ATK','ATK','ATK','ATK']
df1['day'] = [3, 4, 5, 6, 2, 3, 4, 5, 1, 3, 5, 7, 1, 2, 6, 8, 2, 4, 6, 8]
df1

I have tried to use the logical function of any but I could not make it work, it returns only True or False to me instead of a filtered dataframe.

Upvotes: 4

Views: 2198

Answers (2)

EdChum
EdChum

Reputation: 393913

IIUC you want:

In[116]:
df_grouped = df.groupby(['cc', 'odd', 'tree1', 'tree2']).filter(lambda x: (x['day'] > 4).any())
df_grouped

Out[116]: 
    cc   odd tree1 tree2  day
4   BB  3435   SAP   ATK    3
5   BB  3435   SAP   ATK    4
6   BB  3435   SAP   ATK    5
7   BB  3435   SAP   ATK    6
8   DD  3434   ASP   ATK    2
9   DD  3434   ASP   ATK    3
10  DD  3434   ASP   ATK    4
11  DD  3434   ASP   ATK    5
12  DD  3435   SAP   ATK    1
13  DD  3435   SAP   ATK    3
14  DD  3435   SAP   ATK    5
15  DD  3435   SAP   ATK    7
16  ZZ  3434   ASP   ATK    1
17  ZZ  3434   ASP   ATK    2
18  ZZ  3434   ASP   ATK    6
19  ZZ  3434   ASP   ATK    8
20  ZZ  3435   SAP   ATK    2
21  ZZ  3435   SAP   ATK    4
22  ZZ  3435   SAP   ATK    6
23  ZZ  3435   SAP   ATK    8

So this will filter out the groups where within the group none of the 'day' values are greater than 4

timings:

%timeit df[df.day.gt(4).groupby([df.cc, df.odd, df.tree1, df.tree2]).transform('any')]
%timeit df.groupby(['cc', 'odd', 'tree1', 'tree2']).filter(lambda x: (x['day'] > 4).any())
%timeit df[df.assign(key=df.day > 4).groupby(['cc', 'odd', 'tree1', 'tree2']).key.transform('any')]
100 loops, best of 3: 5.9 ms per loop
100 loops, best of 3: 5.42 ms per loop
100 loops, best of 3: 3.62 ms per loop

So @coldspeed's first method is the fastest here

Upvotes: 2

cs95
cs95

Reputation: 402263

Now that I've understood what you want, let's try something like transform + any:

df[df.assign(key=df.day > 4)
     .groupby(['cc', 'odd', 'tree1', 'tree2']).key.transform('any')
]

Or,

df[df.day.gt(4).groupby([df.cc, df.odd, df.tree1, df.tree2]).transform('any')]


    cc   odd tree1 tree2  day
4   BB  3435   SAP   ATK    3
5   BB  3435   SAP   ATK    4
6   BB  3435   SAP   ATK    5
7   BB  3435   SAP   ATK    6
8   DD  3434   ASP   ATK    2
9   DD  3434   ASP   ATK    3
10  DD  3434   ASP   ATK    4
11  DD  3434   ASP   ATK    5
12  DD  3435   SAP   ATK    1
13  DD  3435   SAP   ATK    3
14  DD  3435   SAP   ATK    5
15  DD  3435   SAP   ATK    7
16  ZZ  3434   ASP   ATK    1
17  ZZ  3434   ASP   ATK    2
18  ZZ  3434   ASP   ATK    6
19  ZZ  3434   ASP   ATK    8
20  ZZ  3435   SAP   ATK    2
21  ZZ  3435   SAP   ATK    4
22  ZZ  3435   SAP   ATK    6
23  ZZ  3435   SAP   ATK    8

Upvotes: 4

Related Questions