Выборочная группа, np.where Pandas / Python запрос

Question

df_have

A  B C
1  1  10
1  1  10
1  2  5
1  3  6
1  4  7

df_want

A  B C    D
1  1  10  20
1  1  10  20
1  2  5   5
1  3  6   6
1  4  7   7

Попытка сгруппировать по столбцу A, B, только если столбец B равен 1 или 2. Если B равен 1 или 2, то вывести столбец Dкак СУММА столбца C. Иначе оставьте установленным столбец D = столбец C.

Пример кода дает ошибку:

df_want['D']=np.where((df_want['B'].isin([1,2]), 
             df_want['A','B'].map(df_want.groupby(['A','B'])['C'].sum()), 
             df_want['C'])

Проблема возникает здесь: df_want ['A', 'B'] .map .. если я только поставлю df_want.A.map, то код выполняется, но вывод неправильный.Мне нужно это только для отображения суммы, если столбец B равен 1 или 2

Answer 1

Вы можете использовать loc

df['D'] = df['C'].copy()
df.loc[df['B'].isin([1,2]), 'D'] = df.loc[df['B'].isin([1,2])].groupby(['A', 'B']).C.transform('sum')


    A   B   C   D
0   1   1   10  20
1   1   1   10  20
2   1   2   5   5
3   1   3   6   6
4   1   4   7   7

Answer 2

Измените sum на transform('sum')

np.where(df_want['B'].isin([1,2]), 
             df_want.groupby(['A','B'])['C'].transform('sum'), df_want['C'])