groupby filter for accounts where monthly balance are all negative

user_4098013 0
Viewed 52

Here is my sample data -

import pandas as pd

df = pd.DataFrame({'Account': ['A1', 'A1', 'A1', 'A1', 'A2', 'A2', 'A2', 'A2', 'A3', 'A3', 'A3', 'A3'],
                   'Date': ['D1', 'D2', 'D3', 'D4', 'D1', 'D2', 'D3', 'D4', 'D1', 'D2', 'D3', 'D4'],
                   'Month': ['M1', 'M1', 'M2', 'M2', 'M1', 'M1', 'M2', 'M2', 'M1', 'M1', 'M2', 'M2'],
                   'Balance': [133, 321, 123, 234, 345, 344, 456, -765, -123, -312, 111, -766]})

print(df)

   Account Date Month  Balance
0       A1   D1    M1      133
1       A1   D2    M1      321
2       A1   D3    M2      123
3       A1   D4    M2      234
4       A2   D1    M1      345
5       A2   D2    M1      344
6       A2   D3    M2      456
7       A2   D4    M2     -765
8       A3   D1    M1     -123
9       A3   D2    M1     -312
10      A3   D3    M2      111
11      A3   D4    M2     -766

I want to filter out accounts from this df where all monthly balances for an account is less than 0. In the example, only A3 should get filtered out as all monthly balances for that account is less than 0.

Here is what I tried which does the job -

eod_gr = (df.groupby(['Account', 'Month'])['Balance'].mean() < 0).reset_index()
tmp = eod_gr.groupby('Account')['Balance'].all().to_dict()
accounts = [i for i in tmp if not tmp[i]]
print(accounts)
df_eod = df[df['Account'].isin(accounts)]
print(df_eod)

The problem is I am working on a real time use case where speed is paramount. I need a way to optimize this query to attain the same results.

Expected Output

  Account Date Month  Balance
0      A1   D1    M1      133
1      A1   D2    M1      321
2      A1   D3    M2      123
3      A1   D4    M2      234
4      A2   D1    M1      345
5      A2   D2    M1      344
6      A2   D3    M2      456
7      A2   D4    M2     -765
pandas-groupby pandas python
1 Answers

I think here should be improved performance omit by second groupby, here is used GroupBy.transform with 'invert' mask from < to >= for get all accounts:

mask = df.groupby(['Account', 'Month'])['Balance'].transform('mean') >= 0
accounts = df.loc[mask, 'Account']
df_eod = df[df['Account'].isin(accounts)]
print(df_eod)
  Account Date Month  Balance
0      A1   D1    M1      133
1      A1   D2    M1      321
2      A1   D3    M2      123
3      A1   D4    M2      234
4      A2   D1    M1      345
5      A2   D2    M1      344
6      A2   D3    M2      456
7      A2   D4    M2     -765

Another idea with aggregation mean and filter MultiIndex Series with Index.get_level_values:

s = df.groupby(['Account', 'Month'])['Balance'].mean()
accounts = s.index[s >= 0].get_level_values(0)
df_eod = df[df['Account'].isin(accounts)]
user_29010020
Related