迭代DataFrame列对比相邻值生成新列时出现KeyError及运行慢如何解决

问题根因 #

• 排序未生效：sort_values()默认返回新的DataFrame，原有代码没有接收返回值也没有加inplace=True，排序操作等于没执行
• 索引越界：遍历到最后一行时i+1超出索引范围触发KeyError，外层try-except会直接终止循环，还会吞掉其他潜在异常
• 逻辑风险+效率低下：如果DataFrame索引不是连续整数，i+1根本不是下一行的索引，会得到完全错误的结果；且Python层的for循环遍历在数据量大时速度远低于pandas原生向量化运算

最优实现方案 #

直接用pandas内置的shift()方法做向量化对比，无需循环，无越界问题，性能提升百倍以上：

%%time
# 排序同时重置为连续索引，避免原索引乱序带来的逻辑问题
final_df = final_df.sort_values(by=['MAN'], ignore_index=True)

# 对比当前行与下一行的MAN值，相等赋值0、不等赋值1，最后一行默认赋值0
final_df["MAN_ID"] = (final_df["MAN"] != final_df["MAN"].shift(-1)).astype("int")

# 如果需要最后一行赋值为1，可追加下面这行代码
# final_df.loc[final_df.index[-1], "MAN_ID"] = 1

逻辑说明 #

• shift(-1)会将MAN列整体向上偏移1行，当前行调用shift(-1)的返回值就是下一行的MAN值
• 布尔值转int时False对应0、True对应1，刚好匹配需求的赋值规则
• 最后一行的shift(-1)返回值为NaN，和原值对比结果为False，默认赋值为0，可按需单独调整最后一行的取值

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