首页 > 解决方案 > 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

标签: pythonpandaspandas-groupby

解决方案


我认为这里应该通过第二个 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)]

推荐阅读