如何用groupby从DataFrame多列返回指定日期值？

示例数据构造 #

先创建符合需求的测试DataFrame：

import pandas as pd
import numpy as np

data = {
    'date': ['2/28/2023', '3/1/2023', '3/2/2023', '3/3/2023', '3/4/2023'],
    'column_a': [0, 0, 0, 1, 1],
    'column_b': [0, 0, 1, 1, 0],
    'column_c': [1, 1, 0, 0, 0]
}
df = pd.DataFrame(data)
df['date'] = pd.to_datetime(df['date'])

预期效果：第三行（3/2/2023）的return_date为3/1/2023（column_c连续1的最后日期），其余行填NaN。

实现步骤 #

1. 转换为长格式数据 #

将多列的1/0值转换为长格式，只保留值为1的行，方便后续分组处理：

long_df = df.melt(id_vars='date', var_name='column', value_name='flag')
long_df = long_df[long_df['flag'] == 1]

2. 标记连续1的分组块 #

通过比较当前行与上一行的列名、日期，判断是否为新的连续1块起始，生成唯一的block_id作为分组键：

# 标记新块的起始：列名变化 或 日期不连续
long_df['is_new_block'] = (long_df['column'] != long_df['column'].shift()) | \
                          (long_df['date'] != long_df['date'].shift() + pd.Timedelta(days=1))
long_df['block_id'] = long_df['is_new_block'].cumsum()

3. 计算每个块的起止日期 #

用groupby按block_id聚合，得到每个连续1块的最早（起始）和最晚（结束）日期：

# 每个块的结束日期
block_end = long_df.groupby('block_id').agg(
    end_date=('date', 'max'),
    column=('column', 'first')
).reset_index()

# 每个块的起始日期
block_start = long_df.groupby('block_id').agg(
    start_date=('date', 'min'),
    column=('column', 'first')
).reset_index()

4. 关联前后块的日期 #

将前一个块的结束日期，关联到后一个块的起始记录中：

# 给每个块添加上一个块的结束日期
block_start['prev_block_end'] = block_start['block_id'].shift().map(block_end.set_index('block_id')['end_date'])

5. 映射回原DataFrame生成目标列 #

创建起始日期+列名到前块结束日期的映射，遍历原df找到新块起始行并赋值：

# 构建映射字典
date_col_map = block_start.set_index(['start_date', 'column'])['prev_block_end'].to_dict()

# 初始化目标列
df['return_date'] = np.nan

# 遍历填充目标列
for idx, row in df.iterrows():
    # 找出当前行值为1的列
    active_cols = [col for col in df.columns if col.startswith('column_') and row[col] == 1]
    for col in active_cols:
        # 判断是否是该列连续1的起始行
        if idx == 0 or df.loc[idx-1, col] == 0:
            # 从映射中获取前块结束日期
            prev_date = date_col_map.get((row['date'], col))
            if prev_date is not None:
                df.loc[idx, 'return_date'] = prev_date

groupby的核心作用 #

这里groupby('block_id')是处理多列连续序列的关键：

• 把分散在不同列的连续1序列，统一聚合为独立的分组块
• 方便快速计算每个序列的时间范围（起止日期）
• 后续可以通过分组ID的顺序，关联前后序列的时间关系

运行后得到的df如下：

date  column_a  column_b  column_c return_date
0 2023-02-28         0         0         1         NaT
1 2023-03-01         0         0         1         NaT
2 2023-03-02         0         1         0  2023-03-01
3 2023-03-03         1         1         0         NaT
4 2023-03-04         1         0         0         NaT

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