如何计算DataFrame两列差值并按条件标记移库行与数量

原始数据 #

原始DataFrame如下：

import pandas as pd

df = pd.DataFrame({
    'Group': ['A', 'A', 'A', 'B', 'B', 'C', 'C', 'C', 'D'],
    'Required': [10, 10, 10, 13, 13, 8, 8, 8, 16],
    'stock': [5, 8, 7, 6, 5, 4, 5, 8, pd.NA]
})

表格展示：

需求说明 #

按Group分组，每组对应固定Required值（A:10、B:13、C:8、D:16），需完成：

• 标记每行是否需要移库（flag列：yes/no）
• 计算每行的移库数量（to_move列）
  规则细节：
• 若stock为NaN，直接标记flag=no，to_move=NaN
• 对每组非NaN行按顺序累加stock：
  • 累加未达Required时，该行flag=yes，to_move取当前stock值
  • 累加值首次≥Required时，该行flag=yes，to_move取Required减去之前的累计值；后续行flag=no，to_move=0
  • 若所有行累加后仍未达Required，则所有行flag=yes，to_move取各自stock值

实现代码 #

import pandas as pd

# 初始化数据
df = pd.DataFrame({
    'Group': ['A', 'A', 'A', 'B', 'B', 'C', 'C', 'C', 'D'],
    'Required': [10, 10, 10, 13, 13, 8, 8, 8, 16],
    'stock': [5, 8, 7, 6, 5, 4, 5, 8, pd.NA]
})

# 初始化列值
df['stock'] = df['stock'].astype(float)
df['to_move'] = df['stock'].copy()
df['flag'] = 'yes'

# 定义分组处理函数
def process_group(g):
    required = g['Required'].iloc[0]
    non_nan_rows = g[~g['stock'].isna()]
    
    if len(non_nan_rows) == 0:
        g['flag'] = 'no'
        g['to_move'] = pd.NA
        return g
    
    # 计算累计库存
    cum_stock = non_nan_rows['stock'].cumsum()
    # 找到首次满足需求的行索引
    first_meet_idx = cum_stock[cum_stock >= required].index.min()
    
    if pd.notna(first_meet_idx):
        # 计算之前的累计库存
        prev_total = cum_stock.loc[cum_stock.index < first_meet_idx].sum() if len(cum_stock.index < first_meet_idx) > 0 else 0
        # 更新当前行的移库数量
        g.loc[first_meet_idx, 'to_move'] = required - prev_total
        # 标记后续行无需移库
        g.loc[g.index > first_meet_idx, 'flag'] = 'no'
        g.loc[g.index > first_meet_idx, 'to_move'] = 0.0
    
    return g

# 分组应用处理逻辑
df = df.groupby('Group', group_keys=False).apply(process_group)
# 调整列顺序
df = df[['Group', 'Required', 'stock', 'to_move', 'flag']]
print(df)

输出结果 #

运行代码后得到目标结果：

Group  Required  stock  to_move flag
0     A        10    5.0      5.0  yes
1     A        10    8.0      5.0  yes
2     A        10    7.0      0.0   no
3     B        13    6.0      6.0  yes
4     B        13    5.0      5.0  yes
5     C         8    4.0      4.0  yes
6     C         8    5.0      4.0  yes
7     C         8    8.0      0.0   no
8     D        16    NaN      NaN   no

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