Pandas中基于日期条件的列计算异常问题求助

原始数据与需求 #

构造的DataFrame #

import pandas as pd

data = {
    'Month': ['2023-11','2023-12','2024-01','2024-02','2024-03','2024-04','2024-05','2024-06'], 
    'EU': [30, 40, 29, 25, 27, 27, 37, 29], 
    'Japan': [310, 60, 60, 60, 60, 50, 350, 49], 
    'MTO': [400, 0, 0, 0, 0, 0, 0, 0], 
    'ROW': [2206, 1794, 3185, 1683, 2281, 2240, 2226, 2418], 
    'Month_date': ['2023-11-01', '2023-12-01', '2024-01-01', '2024-02-01', '2024-03-01', '2024-04-01', '2024-05-01', '2024-06-01']
}

forec_pivot2 = pd.DataFrame(data)
forec_pivot2["Month_date"] = pd.to_datetime(forec_pivot2["Month_date"])

需求 #

根据Month_date与阈值2024-02-01的比较结果，计算新列new_col：

• 当Month_date < 2024-02-01时，取值为Japan + MTO
• 当Month_date >= 2024-02-01时，取值为Japan + MTO + ROW

用户尝试的代码及问题 #

用户尝试用循环实现，但出现语法错误、索引越界或逻辑错误：

第一次尝试代码 #

import datetime
d1 = pd.to_datetime('2024-02-01')

list_0=[]
list_1=[]

    for i in range(8):
        if forec_pivot2.loc[i,'Month_date'] < d1:
            result=forec_pivot2.loc[i,'Japan']+forec_pivot2.loc[i,'MTO']
            list_0.append(result)
        else:
            result2=forec_pivot2.loc[i,'Japan']+forec_pivot2.loc[i,'MTO']+forec_pivot2.loc[i,'ROW']
            list_1.append(result2)

final_list = list_0 + list_1
forec_pivot2['new_col'] = final_list

问题：for循环前存在多余缩进，触发语法错误；且循环遍历的方式不符合Pandas的向量化设计，易出现索引匹配问题。

第二次尝试代码 #

for i in range(36):
    if forec_pivot2.loc[i,'Month_date'] < d1:
        result=forec_pivot2.loc[i,'Japan']+forec_pivot2.loc[i,'MTO']
        list_0.append(result)
    else:
        result2=forec_pivot2.loc[i,'Japan']+forec_pivot2.loc[i,'MTO']+forec_pivot2.loc[i,'ROW']
        list_1.append(result)

问题：DataFrame仅8行，range(36)会导致loc[i]超出索引范围，触发KeyError；且else分支错误地将result2写成result，导致逻辑错误。

正确解决方案 #

推荐使用Pandas的向量化操作，高效且避免索引问题：

方案1：使用np.where实现条件赋值 #

import numpy as np

d1 = pd.to_datetime('2024-02-01')
forec_pivot2['new_col'] = np.where(
    forec_pivot2['Month_date'] < d1,
    forec_pivot2['Japan'] + forec_pivot2['MTO'],
    forec_pivot2['Japan'] + forec_pivot2['MTO'] + forec_pivot2['ROW']
)

方案2：使用loc索引分步赋值 #

d1 = pd.to_datetime('2024-02-01')
# 先给所有行赋值为默认的Japan+MTO+ROW
forec_pivot2['new_col'] = forec_pivot2['Japan'] + forec_pivot2['MTO'] + forec_pivot2['ROW']
# 对符合条件的行重新赋值
forec_pivot2.loc[forec_pivot2['Month_date'] < d1, 'new_col'] = forec_pivot2['Japan'] + forec_pivot2['MTO']

执行后，forec_pivot2的new_col列会得到正确结果：

内容的提问来源于stack exchange，提问作者Guillermo Nieva Sánchez