如何让np.where在新1值分段出现时使column2值自动递增？

问题场景 #

现有数据集的column1列仅包含0和1，数据如下：

column1:
1
1
1
1
0
0
0
1
1
0
1
1

需求是创建新列column2，对column1中的连续1分段进行标记：每出现一个新的连续1段，column2的值递增1；0对应的位置保持为0。

当前使用np.where()的代码无法实现需求：

n = 1
df['column2'] = np.where(df['column1'] == 1 ,  n + 1  , 0)

得到的错误结果：

column1: column2:
1            1
1            1
1            1
1            1
0            0
0            0
0            0 
1            1
1            1
0            0
1            1
1            1

期望的正确结果：

column1: column2:
1            1
1            1
1            1
1            1
0            0
0            0
0            0 
1            2
1            2
0            0
1            3
1            3

正确实现代码 #

核心思路是识别每个连续1段的起始位置，通过累加起始位置的数量得到分段编号，再与原列相乘保留0的位置：

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame({'column1': [1,1,1,1,0,0,0,1,1,0,1,1]})

# 识别连续1段的起始行（前一行是0，或当前是第一行且值为1）
segment_starts = (df['column1'] == 1) & (df['column1'].shift(1).fillna(0) != 1)
# 累加起始点得到分段编号，与column1相乘确保0的位置为0
df['column2'] = segment_starts.cumsum() * df['column1']

print(df)

代码解释 #

• df['column1'].shift(1).fillna(0)：将column1向下偏移一行，并用0填充第一行的空值，用于对比当前行与前一行的值
• segment_starts：筛选出所有连续1段的第一个元素，这些位置会被标记为True
• segment_starts.cumsum()：对True值累加，得到每个分段的递增编号
• 最后乘以df['column1']：让原列中0的位置保持为0，1的位置对应正确的分段编号

运行代码后即可得到期望的结果。

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