如何在Pandas中按条件为各ID补充行至day_5？

现有如下Pandas DataFrame：

import pandas as pd
df = pd.DataFrame([["X","day_2"],["Y","day_4"],["Z","day_3"]],columns=["id","day"])

对应表格：

需求：针对每个id，从其day列指定日期的次日开始，补充行直至达到day_5。具体示例：

• ID为X时，补充day_3至day_5，共3行；
• ID为Y时，仅补充day_5这1行；
• ID为Z时，补充day_4和day_5，共2行。

预期输出的DataFrame代码：

df_result = pd.DataFrame([["X","day_3"],["X","day_4"],["X","day_5"],["Y","day_5"],["Z","day_4"],["Z","day_5"]],columns=["id","day"])

对应表格：

这里提供两种实用实现方式：

方法一：逐行生成（逻辑直观） #

1. 提取day列中的数字部分，转为整数方便计算：

df['day_num'] = df['day'].str.extract(r'(\d+)').astype(int)

2. 定义函数生成每个id需要补充的日期序列，再合并结果：

def generate_supplement_days(row):
    start_day = row['day_num'] + 1
    if start_day > 5:
        return pd.DataFrame()
    # 生成从start_day到5的day字符串
    supplement_days = [f'day_{i}' for i in range(start_day, 6)]
    return pd.DataFrame({'id': row['id'], 'day': supplement_days})

# 对每行应用函数并合并所有结果
result = pd.concat([generate_supplement_days(row) for _, row in df.iterrows()], ignore_index=True)

方法二：向量化操作（高效适合大数据） #

通过范围生成+展开的方式避免循环：

# 提取day列末尾的数字并转为整数
df['day_num'] = df['day'].str[-1].astype(int)
# 定义每个id的日期范围（从次日到day_5）
df['date_range'] = df.apply(lambda x: range(x['day_num']+1, 6), axis=1)
# 展开范围并生成对应的day列
result = df.explode('date_range')[['id']].assign(day=lambda x: 'day_' + x['date_range'].astype(str))

两种方法最终得到的result均为预期输出。

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