Группирование по многоуровневой категории и возврату n-наибольшей суммы в каждой категории (n отличается для каждой категории)

Question

У меня есть таблица данных pandas (df), представляющая ежемесячные расходы разных людей. Первый столбец в кадре данных относится к идентификатору лица, второй столбец относится к категории расходов, а третий столбец относится к потраченной сумме. Смотрите пример таблицы ниже:

d = {'PersonID': ['A','A','A','A','A','A','A','A','B','B','B','B','B','B'], 'Category': ['Food','Food','Food','Food','Travel','Travel','Travel','Travel','Food','Food','Food','Travel','Travel','Travel'], 'Expenditure':[10,15,5,20,500,100,1000,2000,10,30,10,800,1000,400]}
df = pd.DataFrame(data=d)

Для каждого человека я хотел бы получить сумму ТРИ крупнейших расходов в категории «Пища» и сумму ДВУ крупнейших расходов в категории «Путешествия».

Для приведенной выше таблицы я хочу следующую таблицу:

Я пытаюсь использовать следующий код, но проблема в том, что я не могу указать разные N-самые большие расходы в разных категориях.

df.groupby(['PersonID','Category'])['Expenditure'].nlargest(2).sum(level=0)

Answer 1

Способ сделать это состоит в том, чтобы сначала разбить ваш фрейм данных по категориям, а затем сгруппировать суммы и объединить результаты вместе:

pd.concat([
df.query('Category == "Food"').groupby(['PersonID','Category'])['Expenditure'].nlargest(3).sum(level=[0,1]),
df.query('Category == "Travel"').groupby(['PersonID','Category'])['Expenditure'].nlargest(2).sum(level=[0,1])
])

Выход:

PersonID  Category
A         Food          45
B         Food          50
A         Travel      3000
B         Travel      1800
Name: Expenditure, dtype: int64

---

Использование словаря и понимания списка:

d = {'Food':2,
     'Travel':3}

pd.concat([df[df['Category'] == c].groupby(['PersonID','Category'])['Expenditure'].nlargest(n).sum(level=[0,1]) for c,n in d.items()])