Selecting rows based on sum over multiindex in Pandas-CodePudding

import pandas as pd
import numpy as np

np.random.seed(365)
rows = 100
data = {'Month': np.random.choice(['2014-01', '2014-02', '2014-03', '2014-04'], size=rows),
        'Code': np.random.choice(['A', 'B', 'C'], size=rows),
        'ColA': np.random.randint(5, 125, size=rows),
        'ColB': np.random.randint(0, 51, size=rows),}
df = pd.DataFrame(data)

df = df[((~((df.Code=='A')&(df.Month=='2014-04')))&(~((df.Code=='C')&(df.Month=='2014-03'))))]
dfg = df.groupby(['Code', 'Month']).sum()

Above gives my dataframe. I want to select only those entries which have sum (of ColA) over 1000 when this sum is performed over level[0]

dfg.ColA.sum(level=[0])

dfg[dfg.ColA.sum(level=[0])>1000]

Above one throw an error? Expected output is :

        ColA  ColB
Code Month              
B    2014-01   477   300
     2014-02   591   167
     2014-03   522   192
     2014-04   367   169
C    2014-01   412   180
     2014-02   275   205
     2014-04   901   309

CodePudding user response：

You need to use groupby transform to broadcast the sum values across level=0 index

dfg[dfg.groupby(level=0)['ColA'].transform('sum').gt(1000)]

              ColA  ColB
Code Month              
B    2014-01   477   300
     2014-02   591   167
     2014-03   522   192
     2014-04   367   169
C    2014-01   412   180
     2014-02   275   205
     2014-04   901   309

CodePudding user response：

another way to do the same:

groups = [g for _,g in df.groupby('Code') if g.ColA.sum()>1000]
pd.concat(groups).groupby(['Code', 'Month']).sum()
'''
              ColA  ColB
Code Month              
B    2014-01   477   300
     2014-02   591   167
     2014-03   522   192
     2014-04   367   169
C    2014-01   412   180
     2014-02   275   205
     2014-04   901   309