按组将Pandas DataFrame末尾零值填充为前一个非零值+1

我之前解决过类似问题，但实际场景更复杂。现有如下Pandas DataFrame：

LOT_ID           STEP_NUMBER
0     4258556          0.0
1     4258556          0.0
2     4258556          1.0
3     4258556          1.0
4     4258556          2.0
5     4258556          2.0
6     4258556          3.0
7     4258556          3.0
8     4258556          4.0
9     4258556          5.0
10    4258556          5.0
11    4258556          6.0
12    4258556          6.0
13    4258556          7.0
14    4258556          7.0
15    4258556          0.0
16    4261512          0.0
17    4261512          1.0
18    4261512          1.0
19    4261512          2.0
20    4261512          2.0
21    4261512          3.0
22    4261512          3.0
23    4261512          4.0
24    4261512          5.0
25    4261512          5.0
26    4261512          6.0
27    4261512          6.0
28    4261512          7.0
29    4261512          7.0
30    4261512          0.0
31    4265282          0.0
32    4265282          1.0
33    4265282          1.0
34    4265282          2.0
35    4265282          2.0
36    4265282          3.0
37    4265282          3.0
38    4265282          4.0
39    4265282          4.0
40    4265282          5.0
41    4265282          5.0
42    4265282          6.0
43    4265282          7.0
44    4265282          7.0
45    4265282          0.0
46    4265706          0.0
47    4265706          1.0
48    4265706          1.0

需要按LOT_ID分组，将每组STEP_NUMBER列末尾的零值替换为该组最后一个非零值加1，最终结果如下：

LOT_ID           STEP_NUMBER
0     4258556          0.0
1     4258556          0.0
2     4258556          1.0
3     4258556          1.0
4     4258556          2.0
5     4258556          2.0
6     4258556          3.0
7     4258556          3.0
8     4258556          4.0
9     4258556          5.0
10    4258556          5.0
11    4258556          6.0
12    4258556          6.0
13    4258556          7.0
14    4258556          7.0
15    4258556          8.0
16    4261512          0.0
17    4261512          1.0
18    4261512          1.0
19    4261512          2.0
20    4261512          2.0
21    4261512          3.0
22    4261512          3.0
23    4261512          4.0
24    4261512          5.0
25    4261512          5.0
26    4261512          6.0
27    4261512          6.0
28    4261512          7.0
29    4261512          7.0
30    4261512          8.0
31    4265282          0.0
32    4265282          1.0
33    4265282          1.0
34    4265282          2.0
35    4265282          2.0
36    4265282          3.0
37    4265282          3.0
38    4265282          4.0
39    4265282          4.0
40    4265282          5.0
41    4265282          5.0
42    4265282          6.0
43    4265282          7.0
44    4265282          7.0
45    4265282          8.0
46    4265706          0.0
47    4265706          1.0
48    4265706          1.0

之前的方案只针对整列末尾的零值替换，现在需要处理多分组的情况。

通过groupby结合自定义函数即可实现需求，以下是两种适配不同场景的实现方式：

场景1：分组末尾仅单个零值 #

如果每组末尾只有一个零值需要替换，用以下代码：

import pandas as pd

def fix_single_trailing_zero(group):
    # 获取组内最后一个非零的STEP_NUMBER值
    non_zero_vals = group[group['STEP_NUMBER'] != 0]['STEP_NUMBER']
    if not non_zero_vals.empty and group.iloc[-1]['STEP_NUMBER'] == 0:
        group.iloc[-1, group.columns.get_loc('STEP_NUMBER')] = non_zero_vals.iloc[-1] + 1
    return group

# 应用分组处理
df = df.groupby('LOT_ID', group_keys=False).apply(fix_single_trailing_zero)

场景2：分组末尾有多个连续零值 #

如果分组末尾可能存在多个连续零值，用以下代码批量替换：

import pandas as pd

def fix_multiple_trailing_zero(group):
    non_zero_vals = group[group['STEP_NUMBER'] != 0]['STEP_NUMBER']
    if non_zero_vals.empty:
        return group
    
    last_non_zero = non_zero_vals.iloc[-1]
    # 生成末尾连续零值的掩码：从后往前累积判断是否为零
    trailing_zero_mask = group['STEP_NUMBER'][::-1].cummin().astype(bool)[::-1]
    # 替换符合条件的零值
    group.loc[trailing_zero_mask, 'STEP_NUMBER'] = last_non_zero + 1
    return group

# 应用分组处理
df = df.groupby('LOT_ID', group_keys=False).apply(fix_multiple_trailing_zero)

代码说明 #

• groupby('LOT_ID', group_keys=False)：按LOT_ID分组，group_keys=False避免分组键被添加到结果索引中，保持原数据结构。
• 自定义函数中先筛选出组内非零值，判断是否存在非零值后，再定位需要替换的末尾零值，用最后一个非零值加1完成替换。

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