基于DataFrame现有列创建新列并拆分指定列的技术需求

给定如下Pandas DataFrame：

import pandas as pd

data = {
    'datetime': ['2023-01-02 12:00:00', '2023-01-02 12:00:00', 
                 '2023-01-03 12:00:00', '2023-01-03 12:00:00', 
                 '2023-01-04 12:00:00', '2023-01-04 12:00:00'],
    'day_fetched': ['2023-01-01 12:00:00', '2023-01-02 12:00:00', 
                    '2023-01-02 12:00:00', '2023-01-03 12:00:00', 
                    '2023-01-03 12:00:00', '2023-01-04 12:00:00'],
    'col_a': [100, 120, 140, 160, 200, 430],
    'col_b': [200, 400, 500, 700, 300, 200]
}
df = pd.DataFrame(data)

需求一：创建day_ahead列 #

当datetime与day_fetched的日期部分不同时，day_ahead取值为2；日期部分相同时取值为1。预期中间结果：

datetime            day_fetched         col_a col_b day_ahead
0  2023-01-02 12:00:00 2023-01-01 12:00:00 100  200    2
1  2023-01-02 12:00:00 2023-01-02 12:00:00 120  400    1
2  2023-01-03 12:00:00 2023-01-02 12:00:00 140  500    2
3  2023-01-03 12:00:00 2023-01-03 12:00:00 160  700    1
4  2023-01-04 12:00:00 2023-01-03 12:00:00 200  300    2
5  2023-01-04 12:00:00 2023-01-04 12:00:00 430  200    1

需求二：拆分列 #

基于day_ahead列，将col_a拆分为col_a_1和col_a_2，col_b拆分为col_b_1和col_b_2：

• 当day_ahead=1时，col_a的值放入col_a_1，col_b的值放入col_b_1，对应_2列设为NaN
• 当day_ahead=2时，col_a的值放入col_a_2，col_b的值放入col_b_2，对应_1列设为NaN
  预期最终结果：

datetime            day_fetched         col_a_1 col_a_2 col_b_1 col_b_2 day_ahead
0  2023-01-02 12:00:00 2023-01-01 12:00:00 NaN     200     NaN     200     2
1  2023-01-02 12:00:00 2023-01-02 12:00:00 120     NaN     100     NaN     1
2  2023-01-03 12:00:00 2023-01-02 12:00:00 NaN     500     NaN     200     2
3  2023-01-03 12:00:00 2023-01-03 12:00:00 160     NaN     100     NaN     1
4  2023-01-04 12:00:00 2023-01-03 12:00:00 NaN     300     NaN     200     2
5  2023-01-04 12:00:00 2023-01-04 12:00:00 430     NaN     100     NaN     1

步骤1：处理日期列并创建day_ahead #

先将字符串类型的日期转换为Pandas datetime类型，再通过日期比较生成目标列：

# 转换为datetime类型
df['datetime'] = pd.to_datetime(df['datetime'])
df['day_fetched'] = pd.to_datetime(df['day_fetched'])

# 用矢量化操作生成day_ahead（效率高于apply，适合大数据量）
import numpy as np
df['day_ahead'] = np.where(
    df['datetime'].dt.date == df['day_fetched'].dt.date,
    1,
    2
)

步骤2：拆分列 #

使用where方法根据day_ahead的条件赋值，不符合条件的位置自动填充NaN：

# 拆分col_a
df['col_a_1'] = df['col_a'].where(df['day_ahead'] == 1)
df['col_a_2'] = df['col_a'].where(df['day_ahead'] == 2)

# 拆分col_b
df['col_b_1'] = df['col_b'].where(df['day_ahead'] == 1)
df['col_b_2'] = df['col_b'].where(df['day_ahead'] == 2)

# 可选：调整列顺序匹配预期结果
df = df[['datetime', 'day_fetched', 'col_a_1', 'col_a_2', 'col_b_1', 'col_b_2', 'day_ahead']]

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