如何简化Pandas中查询、分组求和及条件赋值的冗余代码？

需求回顾 #

需要完成三个操作：

• 查询DataFrame中B<1的行
• 按D列分组并对A列求和
• 将不满足B<1条件的行的A列值设为0

简洁实现代码 #

import pandas as pd
import numpy as np

# 生成原始数据（添加随机种子保证结果可复现）
np.random.seed(42)
df = pd.DataFrame(np.random.randn(5,3), columns=list('ABC'))
df['D'] = [1,2,2,1,2]

# 核心逻辑：一步完成分组求和与条件赋值
# 1. 先计算每个D分组下，B<1的行的A列总和
group_a_sum = df[df['B'] < 1].groupby('D')['A'].sum()
# 2. 映射分组总和到原数据，同时对B>=1的行将A设为0
df['A'] = df['D'].map(group_a_sum).where(df['B'] < 1, 0)

# 输出结果
print(df)

代码解释 #

• 分组求和：先筛选出B<1的行，按D分组计算A列的总和，得到每个D对应的聚合结果。
• 映射与条件赋值：通过map将分组总和匹配到原DataFrame的每一行，再用where方法实现条件判断——仅保留B<1行的分组求和值，其余行的A直接设为0。

输出验证 #

运行后输出与原代码的df2结果一致：

A         B         C  D
0  0.000000  2.288548 -0.251334  1
1 -1.225049 -1.486899 -0.576246  2
2 -1.225049 -2.210152  0.409037  2
3 -0.884401  0.124899  1.719387  1
4 -1.225049  0.988187  0.160168  2

内容的提问来源于stack exchange，提问作者Syrupsystem