python - groupby 过滤每月余额均为负数的帐户
问题描述
这是我的示例数据 -
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
我想从中过滤掉一个帐户df
的所有每月余额都小于 0 的帐户。在示例中,只A3
应该被过滤掉,因为该帐户的所有每月余额都小于 0。
这是我尝试过的工作 -
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)
问题是我正在研究一个速度至关重要的实时用例。我需要一种方法来优化此查询以获得相同的结果。
预期产出
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
解决方案
我认为这里应该通过第二个 groupby 来改进性能,这里GroupBy.transform
与'invert'
mask from <
to 一起使用>=
以获得全部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
mean
聚合和过滤MultiIndex Series
的另一个想法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)]