вывод результата суммирования ячеек в одну ячейку

Question

Есть dataframe:

df = pd.DataFrame({'Site': [1, 1, 1, 2, 2, 2, 3, 3, 3],
                   'Cells': [3, 6, 3, 3, 3, 3, 3, 6, 3],
                   'du': ['3101', '3101', '3101', '3000', '3000', '3000', '3101', '3101', '3101'],
                   'TXRX':  ['2T2R', '2T2R', '2T2R', '4T4R', '4T4R', '4T4R','2T2R', '4T4R', '4T4R'],
                   'dlchbw': [10, 15, 10, 20, 20, 15, 10, 15, 10]
                   })

Задача следующая:

если df['du'] = 3101 и df['cells'] < 9 и df['TXRX'] = 2, то ant_bw = 2 * df['dlchbw']
если df['du'] = 3101 и df['cells'] < 9 и df['TXRX'] = 4, то ant_bw = 4 * df['dlchbw']
antBW_tot = cумма ant_bw, сгруппированно по Site

Попытка решения:

import pandas as pd

def ant_bw(row):
    ant_bw2 = 0
    ant_bw4 = 0
    if '3101' in row['du'] and row['Cells'] < 9:            
        if row['TXRX'] == '2T2R':
            ant_bw2 = row['dlchbw'] * 2
        elif row['TXRX'] == '4T4R':
            ant_bw4 = row['dlchbw'] * 4
        ant_res = ant_bw2 + ant_bw4
        return ant_res

df['antBW'] = df.apply(ant_bw, axis=1)
# df_antBW = df.groupby('siteid')[['antBW']].sum().reset_index()
# df_antBW.rename(columns={'antBW': 'antBW_tot'}, inplace=True)

Подскажите, пожалуйста, как внести в существующий dataframe результат суммирования ant_res отдельным столбцом antBW_tot, чтобы dataframe выглядел, например, вот так:

Site    Cells   du  TXRX    dlchbw  ant_bw  antBW_tot
    1   3   3101    2T2R    10      20.0    70
    1   6   3101    2T2R    15      30.0    
    1   3   3101    2T2R    10      20.0
    2   3   3000    4T4R    20      NaN
    2   3   3000    4T4R    20      NaN
    2   3   3000    4T4R    15      NaN
    3   3   3101    2T2R    10      20.0    120
    3   6   3101    4T4R    15      60.0    
    3   3   3101    4T4R    10      40.0

Т.е., чтобы значение суммы стояло в первой строке Site 1, например, а не дублировалось по всем строкам с Site 1.

Answer 1

mask = df['du'].eq('3101') & df['Cells'].lt(9) & df['TXRX'].isin(['2T2R', '4T4R'])
df.loc[mask, 'ant_bw'] = df.loc[mask, 'dlchbw'].mul(df.loc[mask, 'TXRX'].str[0].astype(int))
mask = df['ant_bw'].notna()
df.loc[mask, 'antBW_tot'] = df[mask].groupby('Site')['ant_bw'].transform('sum')
mask = ~mask | ~df.groupby('Site')['ant_bw'].cumcount().eq(0)
df.loc[mask, 'antBW_tot'] = None
df['antBW_tot'] = df['antBW_tot'].fillna('')
print(df)

   Site  Cells    du  TXRX  dlchbw  ant_bw antBW_tot
0     1      3  3101  2T2R      10    20.0      70.0
1     1      6  3101  2T2R      15    30.0          
2     1      3  3101  2T2R      10    20.0          
3     2      3  3000  4T4R      20     NaN          
4     2      3  3000  4T4R      20     NaN          
5     2      3  3000  4T4R      15     NaN          
6     3      3  3101  2T2R      10    20.0     120.0
7     3      6  3101  4T4R      15    60.0          
8     3      3  3101  4T4R      10    40.0

БЛОГ НА HUSL

вывод результата суммирования ячеек в одну ячейку

Ответы (1 шт):