Pandas：提取连续值1的起止日期，求更优实现方案

现有如下Pandas数据集：

import pandas as pd, numpy as np
dates = pd.date_range('01/01/2022', '01/11/2022', freq = 'D')
values = [0,0,1,1,0,0,1,1,1,0,1]
df = pd.DataFrame({'date': dates, 'value': values})

数据集输出：

date       value
0   2022-01-01  0
1   2022-01-02  0
2   2022-01-03  1
3   2022-01-04  1
4   2022-01-05  0
5   2022-01-06  0
6   2022-01-07  1
7   2022-01-08  1
8   2022-01-09  1
9   2022-01-10  0
10  2022-01-11  1

需求 #

将数据转换为包含start和end列的格式：

• start：连续值1的首次出现日期
• end：连续值1的末次出现日期

期望输出：

start        end
2022-01-03  2022-01-04  
2022-01-07  2022-01-09
2022-01-11  NaT

已尝试的实现代码 #

conditions = [
    (df.value == 1) & (df.value.shift(1) == 0),
    (df.value == 1) & (df.value.shift(-1) == 0)]
choices = ['start', 'end']
df['value'] =  np.select(conditions, choices, default=pd.NA)
df = df.dropna()
result = df.pivot(columns='value')

输出结果：

date              
value   end         start
2       NaT         2022-01-03
3       2022-01-04  NaT
6       NaT         2022-01-07
8       2022-01-09  NaT
10      NaT         2022-01-11

更高效的实现方法 #

可以通过分组连续相同值的方式快速实现，步骤如下：

1. 创建分组标识：标记连续相同value的组
2. 筛选出value=1的组
3. 对每个组取date的最小值（start）和最大值（end）

基础实现代码 #

# 创建连续相同值的分组标识
df['group'] = (df['value'] != df['value'].shift()).cumsum()

# 筛选value=1的组，计算每组的start和end
result = df[df['value'] == 1].groupby('group')['date'].agg(start='min', end='max').reset_index(drop=True)

print(result)

输出结果：

start        end
0 2022-01-03 2022-01-04
1 2022-01-07 2022-01-09
2 2022-01-11 2022-01-11

匹配期望格式（单个1的end显示为空） #

如果需要让单个1的行end显示为缺失值，可添加以下处理：

result['end'] = np.where(result['start'] == result['end'], pd.NA, result['end'])
print(result)

输出：

start        end
0 2022-01-03 2022-01-04
1 2022-01-07 2022-01-09
2 2022-01-11        NaT

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