如果一行满足特定条件，则在多索引数据框中选择整个子组

Question

I want to select a subgroup in a multi-index dataframe if one of the rows in that subset meets a condition. This is a simple data frame to explain my problem:

col1=[0,0,0,0,2,4,6,0,0,0,100,200,300,400]
col2=[0,0,0,0,4,6,8,0,0,0,200,900,400, 500]
col3 = ['T','F','F','F','F','F','F','T','F','F','F','F','F', 'T']

d = {'Unit': [1, 1, 1, 1, 2, 2, 2, 3, 4, 5, 6, 6, 6, 6], 
 'Year': [2014, 2015, 2016, 2017, 2015, 2016, 2017, 2017, 2014, 2015, 2014, 2015, 2016, 2017], 'col1' : col1, 'col2' : col2 }
df = pd.DataFrame(data=d)

new_df = df.groupby(['Unit', 'Year']).sum()

new_df['col3'] = (new_df.groupby(level=0, group_keys=False)
                  .apply(lambda x: x.col1/x.col2.shift())
                 )

           col1  col2      col3
Unit Year                      
1    2014     0     0       T
     2015     0     0       F
     2016     0     0       F
     2017     0     0       F
2    2015     2     4       F
     2016     4     6       F
     2017     6     8       F
3    2017     0     0       T
4    2014     0     0       F
5    2015     0     0       F
6    2014   100   200       F
     2015   200   900       F
     2016   300   400       F
     2017   400   500       T

So I would want to select all subgroups which have for one T in col 3

so my output would look like:

           col1  col2      col3
Unit Year                      
1    2014     0     0       T
     2015     0     0       F
     2016     0     0       F
     2017     0     0       F
3    2017     0     0       T
6    2014   100   200       F
     2015   200   900       F
     2016   300   400       F
     2017   400   500       T

Thank you in advance,

Jen

Answer 1

Use:

col1=[0,0,0,0,2,4,6,0,0,0,100,200,300,400]
col2=[0,0,0,0,4,6,8,0,0,0,200,900,400, 500]
col3 = ['T','F','F','F','F','F','F','T','F','F','F','F','F', 'T']

d = {'Unit': [1, 1, 1, 1, 2, 2, 2, 3, 4, 5, 6, 6, 6, 6], 
 'Year': [2014, 2015, 2016, 2017, 2015, 2016, 2017, 2017, 2014, 2015, 2014, 2015, 2016, 2017], 
         'col1' : col1, 'col2' : col2, 'col3' : col3 }
df = pd.DataFrame(data=d)

df = df.set_index(['Unit','Year'])

df = df[df['col3'].eq('T').astype(int).groupby(level=0).transform('sum').eq(1)]
print (df)
           col1  col2 col3
Unit Year                 
1    2014     0     0    T
     2015     0     0    F
     2016     0     0    F
     2017     0     0    F
3    2017     0     0    T
6    2014   100   200    F
     2015   200   900    F
     2016   300   400    F
     2017   400   500    T

Details :

Compare column for equality by Series.eq and cast to integers:

print (df['col3'].eq('T').astype(int))
Unit  Year
1     2014    1
      2015    0
      2016    0
      2017    0
2     2015    0
      2016    0
      2017    0
3     2017    1
4     2014    0
5     2015    0
6     2014    0
      2015    0
      2016    0
      2017    1
Name: col3, dtype: int32

Then count sum per first level with GroupBy.transform for get same size Series :

print (df['col3'].eq('T').astype(int).groupby(level=0).transform('sum'))
Unit  Year
1     2014    1
      2015    1
      2016    1
      2017    1
2     2015    0
      2016    0
      2017    0
3     2017    1
4     2014    0
5     2015    0
6     2014    1
      2015    1
      2016    1
      2017    1
Name: col3, dtype: int32

Compare by 1 and last filter by boolean indexing :

print (df[df['col3'].eq('T').astype(int).groupby(level=0).transform('sum').eq(1)])
           col1  col2 col3
Unit Year                 
1    2014     0     0    T
     2015     0     0    F
     2016     0     0    F
     2017     0     0    F
3    2017     0     0    T
6    2014   100   200    F
     2015   200   900    F
     2016   300   400    F
     2017   400   500    T