pandas groupby條件行和

Question

我有一個如下數據框：

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_1 >= col_2行上對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())

但是，當數據變大（數百萬行）時，這種方法會非常慢，是否有更有效的方法？

Answer 1

只是提前做條件數學。

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