Может ли функция LAMBDA при агрегировании брать условие из другого столбца в python?

Question

Я ищу способ, если он существует, для выполнения агрегации на df, используя только подход lambda, при условии выполнения условия из другого столбца. Вот небольшой микромир проблемы.

df = pd.DataFrame({'ID':[1,1,1,1,2,2],
                   'revenue':[40,55,75,80,35,60],
                   'month':['2012-01-01','2012-02-01','2012-01-01','2012-03-01','2012-02-01','2012-03-01']})
print(df)
   ID       month  revenue
0   1  2012-01-01       40
1   1  2012-02-01       55
2   1  2012-01-01       75
3   1  2012-03-01       80
4   2  2012-02-01       35
5   2  2012-03-01       60

Если вам нужно иметь уникальные месяцы для каждого ID, тогда подойдет следующий код (этот код предназначен только для демонстрации, 'month':'nunique' здесь работает) .

df = df.groupby(['ID']).agg({'month':lambda x:x.nunique()}).reset_index()
print(df)
   ID  month
0   1      3
1   2      2

Но мне нужно подсчитать уникальные months, когда доход был больше 50, взяв две переменные (доход и месяц) в lambda что-то вроде lambda x,y: ....

Я мог бы сделать это как df[df['revenue'] > 50].groupby.(....), но в agg() есть много других столбцов, где это условие не требуется. Итак, существует ли подход, при котором lambda может принимать 2 переменные одновременно ??

Ожидаемый результат:

   ID  month
0   1      3
1   2      1

Answer 1

К сожалению, это не простой / производительный способ, потому что GroupBy.agg обрабатывает каждый столбец отдельно:

Не используйте его, потому что он очень медленный, если большой df или много групп .

def f(x):
   a = df.loc[x.index]
   return a.loc[a['revenue'] > 50, 'month'].nunique()


df1 = df.groupby(['ID']).agg({'month':f}).reset_index()
print(df1)
   ID  month
0   1      3
1   2      1

Таким образом, одним из возможных решений является фильтрация перед использованием GroupBy.apply.