如何基于DataFrame下一行值生成上限列并修复代码缺陷？

需求与现状 #

• 需求：为DataFrame添加「Upper Bound」列，规则如下：
  • 同Item Internal ID分组下，当前行上限值等于下一行的From Quantity减1
  • 每个分组的最后一行默认设为1000
• 现状：原有循环代码仅在数据全局有序时正常运行，无序场景下逻辑失效，且未正确处理分组的最后一行：

for i in range(len(df_p)):
        if( i < (len(df_p)-1)):
            df_p.iloc[i, df_p.columns.get_indexer(['Upper Bound'])] = np.where(((df_p.iloc[i, df_p.columns.get_indexer(['Item Internal ID'])] == df_p.iloc[i+1, df_p.columns.get_indexer(['Item Internal ID'])])), df_p.iloc[i+1, df_p.columns.get_indexer(['From Quantity'])] - 1, 100000 )

优化方案 #

核心思路是按Item Internal ID分组，在组内对From Quantity排序后计算上下限，彻底解决无序数据的适配问题，同时正确处理每组最后一行。

1. 分组排序并获取组内下一行的数值 #

先对数据按Item Internal ID分组，组内按From Quantity升序排序，再通过shift(-1)获取每组内下一行的From Quantity值：

# 按分组排序，重置索引避免原索引干扰
df_p = df_p.sort_values(['Item Internal ID', 'From Quantity']).reset_index(drop=True)

# 分组后获取下一行的From Quantity，组内最后一行会得到NaN
df_p['Next_From_Qty'] = df_p.groupby('Item Internal ID')['From Quantity'].shift(-1)

2. 计算「Upper Bound」列 #

使用np.where判断是否存在下一行数据，分别赋值：

import numpy as np

df_p['Upper Bound'] = np.where(
    df_p['Next_From_Qty'].notna(),  # 存在下一行时
    df_p['Next_From_Qty'] - 1,      # 取下一行值减1
    1000                            # 组内最后一行设为1000
)

# 清理临时列（可选）
df_p = df_p.drop(columns=['Next_From_Qty'])

方案优势 #

• 适配无序数据：通过组内排序，确保每组内的行按From Quantity顺序排列，不受原始数据顺序影响
• 正确处理边界：shift(-1)自动标记组内最后一行，无需手动判断行索引
• 性能提升：用Pandas矢量化操作替代循环，处理大数据量时效率远高于原代码

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