在Pandas DataFrame中按分组规则创建计算列

刚好遇到过类似的场景，我来帮你把这个逻辑落地，先从你提到的思路入手，再给你推荐更高效的Pandas原生写法～

需求回顾 #

给定一个按name连续分组排序的Pandas DataFrame，需要新增计算列calc，规则如下：

• 如果分组的第一个value是0：组内所有行的calc取该组的首个非零值，如果组内全是0则赋值nan
• 如果分组的第一个value非0：组内所有行的calc统一赋值-1

示例输入输出 #

输入DataFrame

import pandas as pd

df = pd.DataFrame({
    'name': ['a', 'a', 'a', 'a', 'b', 'b', 'c', 'c'],
    'value': [0, 0, 6, 8, 0, 0, 5, 7]
})

目标输出

name  value  calc
0    a      0   6.0
1    a      0   6.0
2    a      6   6.0
3    a      8   6.0
4    b      0   NaN
5    b      0   NaN
6    c      5  -1.0
7    c      7  -1.0

方法一：按你的思路实现（生成查找表+遍历） #

这个方法逻辑直观，完全贴合你提到的思路：

1. 先生成每个分组的首个非零值查找表
2. 遍历原DataFrame，根据分组的第一个值判断，从查找表取值或赋值-1

# 步骤1：生成首个非零值查找表
first_non_zero = df[df['value'] != 0].groupby('name')['value'].first()
# 输出结果：
# name
# a    6
# c    5
# Name: value, dtype: int64

# 步骤2：遍历生成calc列
calc_list = []
current_group = None
first_val_in_group = None

for idx, row in df.iterrows():
    # 切换分组时记录当前组的第一个value
    if row['name'] != current_group:
        current_group = row['name']
        first_val_in_group = row['value']
    
    # 根据规则赋值
    if first_val_in_group != 0:
        calc_list.append(-1)
    else:
        # 查找表中有值就取，没有则返回nan
        calc_list.append(first_non_zero.get(current_group, pd.NA))

df['calc'] = calc_list

方法二：更高效的Pandas原生实现（避免遍历） #

遍历的方式在数据量小的时候没问题，但数据量大时效率偏低。推荐用Pandas的groupby+transform/apply的方式，利用向量化操作提升效率，代码也更简洁：

写法1：分组自定义函数

def process_group(group):
    # 获取分组的第一个value
    first_val = group['value'].iloc[0]
    if first_val != 0:
        # 第一个值非零，组内全赋值-1
        return [-1] * len(group)
    else:
        # 找组内首个非零值，没有则返回nan
        non_zero_vals = group[group['value'] != 0]['value']
        target_val = non_zero_vals.iloc[0] if not non_zero_vals.empty else pd.NA
        return [target_val] * len(group)

# 分组处理并展开结果
df['calc'] = df.groupby('name', group_keys=False).apply(process_group).explode()

写法2：先聚合再合并

先把每个分组的关键信息（第一个值、首个非零值）聚合出来，再合并到原表计算：

# 聚合每个分组的关键信息
group_summary = df.groupby('name').agg(
    first_value=('value', 'first'),
    first_non_zero=('value', lambda x: x[x != 0].iloc[0] if any(x != 0) else pd.NA)
)

# 合并到原DataFrame并计算calc列
df = df.merge(group_summary, on='name', how='left')
df['calc'] = df.apply(lambda row: -1 if row['first_value'] != 0 else row['first_non_zero'], axis=1)

# 可选：删除中间辅助列
df = df.drop(['first_value', 'first_non_zero'], axis=1)

两种方法都能得到目标结果，如果你处理的是小数据集，方法一足够直观；如果是大数据量，方法二的效率会高很多。

内容的提问来源于stack exchange，提问作者Krzysztof Słowiński