简体   繁体   中英

Selecting rows based on sum over multiindex in Pandas

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

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

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM