выбрать несколько n-х значений в группировке с условным агрегированием - pandas

Question

у меня есть pd.DataFrame с четырьмя столбцами

    df = pd.DataFrame({'id':[1,1,1,1,1,2,2,2,2] 
                      , 'A':['H','H','E','E','H','E','E','H','H']
                      , 'B':[4,5,2,7,6,1,3,1,0]
                      , 'C':['M','D','M','D','M','M','M','D','D']})

   id  A  B  C
0   1  H  4  M
1   1  H  5  D
2   1  E  2  M
3   1  E  7  D
4   1  H  6  M
5   2  E  1  M
6   2  E  3  M
7   2  H  1  D
8   2  H  0  D

Я хотел бы сгруппировать по идентификатору и получить значение B для n-го (скажем, второго) появления A = ' H 'для каждого идентификатора в agg_B1 и значение B для n-го (скажем, первого) появления C =' M ':

desired output:
   id  agg_B1  agg_B2
0   1      5      4  
1   2      0      1

desired_output = df.groupby('id').agg(
      agg_B1= ('B',lambda x:x[df.loc[x.index].loc[df.A== 'H'][1]])
   ,  agg_B2= ('B',lambda x:x[df.loc[x.index].loc[df.C== 'M'][0]])

    ).reset_index()

TypeError: Indexing a Series with DataFrame is not supported, use the appropriate DataFrame column

Очевидно, я делаю что-то не так с индексацией.

---

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

Answer 1

Ваше решение возможно изменить при необходимости GroupBy.agg:

desired_output = df.groupby('id').agg(
      agg_B1= ('B',lambda x:x[df.loc[x.index, 'A']== 'H'].iat[1]),
      agg_B2= ('B',lambda x:x[df.loc[x.index, 'C']== 'M'].iat[0])
    ).reset_index()


print (desired_output)
   id  agg_B1  agg_B2
0   1       5       4
1   2       0       1

Но если производительность важна, а также не уверен, всегда ли существует второе значение, совпадающее H для первого условия Я предлагаю обрабатывать каждое условие отдельно и в последнюю очередь добавлять к исходным агрегированным значениям:

#some sample aggregations
df0 = df.groupby('id').agg({'B':'sum', 'C':'last'})

df1 = df[df['A'].eq('H')].groupby("id")['B'].nth(1).rename('agg_B1')
df2 = df[df['C'].eq('M')].groupby("id")['B'].first().rename('agg_B2')

desired_output = pd.concat([df0, df1, df2], axis=1)
print (desired_output)
     B  C  agg_B1  agg_B2
id                       
1   24  M       5       4
2    5  D       0       1

EDIT1: при необходимости GroupBy.agg можно проверить, не удалось ли индексировать, а затем добавить недостающее значение:

#for second value in sample working nice
def f1(x):
    try:
        return x[df.loc[x.index, 'A']== 'H'].iat[1]
    except:
        return np.nan

desired_output = df.groupby('id').agg(
      agg_B1= ('B',f1),
      agg_B2= ('B',lambda x:x[df.loc[x.index, 'C']== 'M'].iat[0])
    ).reset_index()

print (desired_output)
   id  agg_B1  agg_B2
0   1       5       4
1   2       0       1

#third value not exist so added missing value NaN
def f1(x):
    try:
        return x[df.loc[x.index, 'A']== 'H'].iat[2]
    except:
        return np.nan

desired_output = df.groupby('id').agg(
      agg_B1= ('B',f1),
      agg_B2= ('B',lambda x:x[df.loc[x.index, 'C']== 'M'].iat[0])
    ).reset_index()

print (desired_output)
   id  agg_B1  agg_B2
0   1     6.0       4
1   2     NaN       1

Что работает аналогично:

df1 = df[df['A'].eq('H')].groupby("id")['B'].nth(2).rename('agg_B1')
df2 = df[df['C'].eq('M')].groupby("id")['B'].first().rename('agg_B2')

desired_output = pd.concat([df1, df2], axis=1)
print (desired_output)
    agg_B1  agg_B2
id                
1      6.0       4
2      NaN       1

Answer 2

Фильтр для строк, где A равно H, затем захватите вторую строку с помощью функции nth :

df.query("A=='H'").groupby("id").nth(1)

    A   B
id      
1   H   5
2   H   0

Python работает с нотацией с нулевым отсчетом, поэтому строка 2 будет nth(1)