Python Dataframe Panda - вычисление взвешенной суммы, если условие соответствует, и группировка результатов

Я хотел бы рассчитать взвешенную сумму на основе условия в столбце criteria . У меня есть следующий Dataframe:

df = pd.DataFrame({'criteria': [2, 2, 2, 3, 4, 5, 5, 9], 
                   'weight': [1, 2, 3, 1, 2, 1, 2, 3], 
                   'value': [0.2, 0.6, 0.1, 0.9, 0.8, 0.7, 0.3, 0.8]})
print(df)

   criteria  weight  value
0         2       1    0.2
1         2       2    0.6
2         2       3    0.1
3         3       1    0.9
4         4       2    0.8
5         5       1    0.7
6         5       2    0.3
7         9       3    0.8

Возможные значения столбца weight : {1, 2, 3}.

Я хотел бы вычислить взвешенную сумму df['weight'] * df['value'] тогда и только тогда, когда есть строка для всех возможных weights и одного и того же значения criteria , а затем сгруппировать их по значениям criteria . Остальное будет NaN .

В моем примере, только строки с df['criteria'] == 2 заполняют это условие.

Таким образом, мы бы имели:

   criteria  weighted_sum
0         2           1.7
1         3           NaN
2         4           NaN
3         5           NaN
4         9           NaN

Как написать это? Большое спасибо!

Всего 1 ответ


Вы можете использовать issubset чтобы проверить, являются ли веса подмножеством группы, а затем взять сумму с помощью series.where :

weights = [1,2,3]
cond = df.groupby('criteria')['weight'].agg(lambda x: set(weights).issubset(set(x)))
(df['weight'].mul(df['value']).groupby(df['criteria']).sum()
                     .where(cond).reset_index(name='weighted_sum'))

   criteria  weighted_sum
0         2           1.7
1         3           NaN
2         4           NaN
3         5           NaN
4         9           NaN