сгруппированная сумма в Pandas / Python с условиями

Question

df_have

ID  AMT     REL_NUM HDR_NUM
3   0.02    2.0      2.0    
4   2.00    2.0      4.0    
5   0.00    1.0      5.0    
1   0.00    5.0      1.0    
2   19.7    1.0      2.0    

df_want

ID  AMT     REL_NUM HDR_NUM CALCULATION
3   0.02    2.0      2.0      (19.7+0.02+2.00)
4   2.00    2.0      4.0      (2.00)
5   0.00    1.0      5.0      (0.00+0.00)
1   0.00    5.0      1.0      (0.00+19.7)
2   19.7    1.0      2.0      (19.7+0.02+2.00)

Попытка создать новый столбец 'CALCULATION', но логика немного хитрая.Вычисление должно быть SUM поля AMT, зависящего от полей ID, REL_NUM и HDR_NUM.

Шаг 1 - Проверьте поле HDR_NUM и получите соответствующее значение AMT, где ID = HDR_NUM. Шаг 2 - Добавьте во все поля AMT, где REL_NUM = HDR_NUM

Для первой строки это будетсуммирование полей AMT для ID = 2,3 и 4.

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

df_want['CALCULATION']=df_have.groupby(['ID','HDR_NUM'])['AMT'].transform('sum')+ ?

Answer 1

Вы можете достичь этого с помощью .map.Для второго вам нужно сгруппировать, чтобы получить сумму в каждом 'REL_NUM'

df['num1'] = df.HDR_NUM.map(df.set_index('ID').AMT)
df['num2'] = df.HDR_NUM.map(df.groupby('REL_NUM').AMT.sum())

df['calculation'] = df.num1.add(df.num2, fill_value=0)

Вывод: (оставляя вспомогательные столбцы для ясности)

   ID    AMT  REL_NUM  HDR_NUM  num1   num2  calculation
0   3   0.02      2.0      2.0  19.7   2.02        21.72
1   4   2.00      2.0      4.0   2.0    NaN         2.00
2   5   0.00      1.0      5.0   0.0   0.00         0.00
3   1   0.00      5.0      1.0   0.0  19.70        19.70
4   2  19.70      1.0      2.0  19.7   2.02        21.72

---

В случае, если вы не хотите дублировать AMT, потому что HDR_NUM == ID == REL_NUM, вы можете суммировать только подмножество для groupby, чтобы не выполнять двойной счет:

df['num1'] = df.HDR_NUM.map(df.set_index('ID').AMT)
df['num2'] = df.HDR_NUM.map(df[df.REL_NUM != df.ID].groupby('REL_NUM').AMT.sum())
df['calculation'] = df.num1.add(df.num2, fill_value=0)

Вывод: (поменялись идентификаторы для первой и последней строки)

   ID    AMT  REL_NUM  HDR_NUM  num1  num2  calculation
0   2   0.02      2.0      2.0  0.02   2.0         2.02
1   4   2.00      2.0      4.0  2.00   NaN         2.00
2   5   0.00      1.0      5.0  0.00   0.0         0.00
3   1   0.00      5.0      1.0  0.00  19.7        19.70
4   3  19.70      1.0      2.0  0.02   2.0         2.02