如何为Pandas DataFrame动态添加重复日期列及循环轮班列

原始DataFrame #

import pandas as pd
import random

data = {
    'Column1': [random.randint(1, 100) for _ in range(9)],
    'Column2': [random.uniform(0, 1) for _ in range(9)],
    'Column3': [chr(random.randint(65, 90)) for _ in range(9)],
    'Column4': [random.choice(['A', 'B', 'C']) for _ in range(9)]
}
df = pd.DataFrame(data)

生成的DataFrame如下：

Column1   Column2 Column3 Column4
0       87  0.208179       M       C
1       85  0.049071       Q       C
2        4  0.474926       X       C
3       35  0.966357       L       B
4       58  0.295134       C       B
5       23  0.633367       R       B
6       87  0.069583       V       B
7       83  0.427594       N       A
8       16  0.592413       R       C

需求说明 #

需要为上述DataFrame新增两列：

• DATE列：从指定start_date开始按顺序生成日期，每个日期重复N次
• SHIFT列：循环使用selected_values列表中的值，列表长度为n

其中N、n、selected_values、start_date均为可调整的动态参数。

实现代码 #

import pandas as pd
import numpy as np

def add_date_shift_columns(df, start_date, N, selected_values):
    # 计算所需日期总数：向上取整总行数与单日期重复次数的比值
    num_dates = (len(df) + N - 1) // N
    # 生成日期序列并重复N次，截取到与原DataFrame匹配的长度
    dates = pd.date_range(start=start_date, periods=num_dates).repeat(N)[:len(df)]
    df['DATE'] = dates.strftime('%Y-%m-%d')  # 可保留datetime类型，此处转为字符串格式
    
    # 循环生成轮班值，截取到与原DataFrame匹配的长度
    shifts = np.tile(selected_values, (len(df) // len(selected_values)) + 1)[:len(df)]
    df['SHIFT'] = shifts
    
    return df

示例验证 #

示例1：N=2、start_date='2022-01-01'、selected_values=['Day','Night'] #

调用代码：

df_result1 = add_date_shift_columns(df.copy(), '2022-01-01', 2, ['Day','Night'])
print(df_result1)

输出结果：

Column1   Column2 Column3 Column4        DATE  SHIFT
0       87  0.208179       M       C  2022-01-01    Day
1       85  0.049071       Q       C  2022-01-01  Night
2        4  0.474926       X       C  2022-01-02    Day
3       35  0.966357       L       B  2022-01-02  Night
4       58  0.295134       C       B  2022-01-03    Day
5       23  0.633367       R       B  2022-01-03  Night
6       87  0.069583       V       B  2022-01-04    Day
7       83  0.427594       N       A  2022-01-04  Night
8       16  0.592413       R       C  2022-01-05    Day

示例2：N=2、start_date='2022-01-01'、selected_values=['Day','Night','Afternoon'] #

调用代码：

df_result2 = add_date_shift_columns(df.copy(), '2022-01-01', 2, ['Day','Night','Afternoon'])
print(df_result2)

输出结果：

Column1   Column2 Column3 Column4        DATE       SHIFT
0       87  0.208179       M       C  2022-01-01         Day
1       85  0.049071       Q       C  2022-01-01       Night
2        4  0.474926       X       C  2022-01-02  Afternoon
3       35  0.966357       L       B  2022-01-02         Day
4       58  0.295134       C       B  2022-01-03       Night
5       23  0.633367       R       B  2022-01-03  Afternoon
6       87  0.069583       V       B  2022-01-04         Day
7       83  0.427594       N       A  2022-01-04       Night
8       16  0.592413       R       C  2022-01-05  Afternoon

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