pandas groupby условная сумма строки - PullRequest
Новая
       20

pandas groupby условная сумма строки

0 голосов
/ 20 марта 2019

У меня есть фрейм данных, как показано ниже: 

df = pd.DataFrame({'col_1': [2,2,2,3,3,3,3],
                   'col_2': [1,2,3,1,2,3,4],
                   'col_3':['A','A','A','B','B','B','B']})
   col_1  col_2 col_3
0      2      1     A
1      2      2     A
2      2      3     A
3      3      1     B
4      3      2     B
5      3      3     B
6      3      4     B

Я хочу суммировать col_2 по строкам с col_1 >= col_2.

Желаемый результат должен быть: 

   col_1  col_2 col_3  col_4
0      2      1     A      3
1      2      2     A      3
2      2      3     A      3
3      3      1     B      6
4      3      2     B      6
5      3      3     B      6
6      3      4     B      6

Я пришел к: 

df.groupby(['col_3']).apply(lambda x: x.loc[x['col_1'] >= x['col_2']]['col_2'].sum())

Это работает, однако, будет очень медленно, когда данные станут большими (миллионы строк), есть ли более эффективный способ?

1 Ответ

0 голосов
/ 20 марта 2019

просто сделай условную математику раньше времени. 

In [46]: df = pd.DataFrame({'col_1': [2,2,2,3,3,3,3],
       :                    'col_2': [1,2,3,1,2,3,4],
       :                    'col_3':['A','A','A','B','B','B','B']})

In [47]: df['cond_val'] = (df.col_1 >= df.col_2) * df.col_2

In [48]: df
Out[48]:
   col_1  col_2 col_3  cond_val
0      2      1     A         1
1      2      2     A         2
2      2      3     A         0
3      3      1     B         1
4      3      2     B         2
5      3      3     B         3
6      3      4     B         0


In [50]: df.groupby('col_3').cond_val.sum()
Out[50]:
col_3
A    3
B    6
Name: cond_val, dtype: int64
...