Pandas新手求助：将含retain语句的SAS代码转换为Pandas代码遇阻

SAS代码逻辑解析 #

这段SAS代码用于从数据集tmp1生成新数据集tmp2，新增group列的规则如下：

• 初始化group起始值为0
• 处理第一行数据时：
  • 若group_v1（tmp1中的一列）等于-1，group设为-1
  • 若group_v1不等于-1，group保持为0
• 从第二行开始，每行group值等于上一行的group值加1

Pandas实现代码 #

新手易懂的循环写法

import pandas as pd

# 复制原DataFrame，避免修改原数据
tmp2 = tmp1.copy()

# 初始化group列为0
tmp2['group'] = 0

# 处理第一行的group值
if tmp2['group_v1'].iloc[0] == -1:
    tmp2.loc[0, 'group'] = -1
else:
    tmp2.loc[0, 'group'] = 0

# 从第二行开始，逐行计算group值
for i in range(1, len(tmp2)):
    tmp2.loc[i, 'group'] = tmp2.loc[i-1, 'group'] + 1

高效无循环写法（适合大数据量）

import pandas as pd

tmp2 = tmp1.copy()

# 确定第一行的group初始值
first_group_val = -1 if tmp2['group_v1'].iloc[0] == -1 else 0

# 构造累加序列：第一行加0，后续每行加1
accumulator = [0] + [1] * (len(tmp2) - 1)
# 计算最终group列
tmp2['group'] = first_group_val + pd.Series(accumulator).cumsum()

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