Как выбрать условные строки с помощью groupby?

Question

Я хочу выбрать строки с групповыми условиями.

import pandas as pd
import numpy as np

dftest = pd.DataFrame({'A':['Feb',np.nan,'Air','Flow','Feb',
                            'Beta','Cat','Feb','Beta','Air'],
                       'B':['s','s','t','s','t','s','t','t','t','t'],
                       'C':[5,4,3,2,1,7,6,5,4,3],
                       'D':[4,np.nan,3,np.nan,2,
                            np.nan,2,3,np.nan,7]})
def filcols3(df,dd):
    if df.iloc[0]['D']==dd:
        return df
dd=4    
grp=dftest.groupby('B').apply(filcols3,dd)

результат grp:

         A  B  C    D
B                   
s 0   Feb  s  5  4.0
  1   NaN  s  4  NaN
  3  Flow  s  2  NaN
  5  Beta  s  7  NaN

это то, что я хочу.

пока я использую следующий код (часть 2)

def filcols3(df,dd):
    if df.iloc[0]['D']<=dd:
        return df
dd=3

результат:

       A    B    C    D
0   NaN  NaN  NaN  NaN
1   NaN  NaN  NaN  NaN
2   Air    t  3.0  3.0
3   NaN  NaN  NaN  NaN
4   Feb    t  1.0  2.0
5   NaN  NaN  NaN  NaN
6   Cat    t  6.0  2.0
7   Feb    t  5.0  3.0
8  Beta    t  4.0  NaN
9   Air    t  3.0  7.0

Я удивлен таким результатом, я хочу получить

      A  B  C    D
2   Air  t  3  3.0
4   Feb  t  1  2.0
6   Cat  t  6  2.0
7   Feb  t  5  3.0
8  Beta  t  4  NaN
9   Air  t  3  7.0

что не так с кодом части 2? как получить желаемый конечный результат?

Answer 1

Может проверить с filter

dftest.groupby('B').filter(lambda x : any(x['D'].head(1)<=3))
Out[538]: 
      A  B  C    D
2   Air  t  3  3.0
4   Feb  t  1  2.0
6   Cat  t  6  2.0
7   Feb  t  5  3.0
8  Beta  t  4  NaN
9   Air  t  3  7.0

Или без groupby drop_duplicates

s=df.drop_duplicates('B').D<=3
df[df.B.isin(df.loc[s.index,'B'][s])]
Out[550]: 
      A  B  C    D
2   Air  t  3  3.0
4   Feb  t  1  2.0
6   Cat  t  6  2.0
7   Feb  t  5  3.0
8  Beta  t  4  NaN
9   Air  t  3  7.0

Answer 2

Поведение

apply здесь немного неинтуитивно, но если идея состоит в том, чтобы отфильтровывать целые группы на основе определенного условия для каждой группы, вы можете использовать GroupBy.transform и получить маску для фильтрации df :

df[df.groupby('B')['D'].transform('first') <= 3]

      A  B  C    D
2  Air   t  3  3.0
4  Feb   t  1  2.0
6  Cat   t  6  2.0
7  Feb   t  5  3.0
8  Beta  t  4 NaN 
9  Air   t  3  7.0

Или, исправляя свой код,

df[df.groupby('B')['D'].transform(lambda x: x.values[0] <= 3)]

      A  B  C    D
2  Air   t  3  3.0
4  Feb   t  1  2.0
6  Cat   t  6  2.0
7  Feb   t  5  3.0
8  Beta  t  4 NaN 
9  Air   t  3  7.0