Python判断日期在指定区间并为Pandas DataFrame计算新增列

实现方案 #

前置准备 #

首先需要统一所有日期字段的类型为datetime，避免字符串比较出现逻辑错误：

import pandas as pd

# 转换日期字段格式
df['r_date'] = pd.to_datetime(df['r_date'])
res_df['date_from'] = pd.to_datetime(res_df['date_from'])
res_df['date_to'] = pd.to_datetime(res_df['date_to'])

方案1：优化版向量化实现（推荐，适合大表） #

避免嵌套循环带来的性能问题，用pandas内置的向量化操作实现：

def calc_new_value(df, res_df):
    # 交叉连接两张表生成所有可能的匹配组合
    cross_df = df.assign(key=1).merge(res_df.assign(key=1), on='key').drop('key', axis=1)
    # 过滤同时满足两个区间条件的匹配项
    match_df = cross_df[
        (cross_df['r_date'].between(cross_df['date_from'], cross_df['date_to'])) &
        (cross_df['num_days'].between(cross_df['days_from'], cross_df['days_to']))
    ]
    # 取每个原表行的第一个匹配系数（如有多个匹配规则，可在此处调整排序逻辑）
    match_df = match_df.groupby(list(df.columns)).head(1)[['koeff']]
    # 合并匹配结果回原表，无匹配的系数默认填1
    df = df.merge(match_df, left_index=True, right_index=True, how='left')
    df['koeff'] = df['koeff'].fillna(1)
    # 计算最终new_value
    df['new_value'] = df['value'] * df['koeff']
    # 删除临时辅助列
    df.drop('koeff', axis=1, inplace=True)
    return df

方案2：基于原有循环的修改版（适合小表，逻辑直观） #

你原有代码存在三个问题：

1. 循环内index变量重名，外层行索引被内层覆盖无法正确赋值
2. 两个表的字段归属搞反，判断逻辑错位
3. 只有条件判断打印，没有值赋值逻辑
   修改后代码如下：

def test_lst(df, res_df):
    # 初始化new_value为原value值
    df['new_value'] = df['value'].copy()
    # 遍历原表每一行
    for df_idx, df_row in df.iterrows():
        # 遍历规则表每一行
        for res_idx, res_row in res_df.iterrows():
            # 判断两个区间是否同时满足
            if (res_row['date_from'] <= df_row['r_date'] <= res_row['date_to']) and \
               (res_row['days_from'] <= df_row['num_days'] <= res_row['days_to']):
                # 匹配到规则则更新值，匹配到第一个规则后跳出循环
                df.loc[df_idx, 'new_value'] = df_row['value'] * res_row['koeff']
                break
    return df

调用示例 #

# 构造示例数据
df = pd.DataFrame({
    'r_date': ['2018-05-13', '2018-06-09', '2018-06-09'],
    'num_days': [3, 9, 12],
    'product': ['CARD', 'CARD', 'AUTO'],
    'value': [11.0, 67.3, 11.0]
})
res_df = pd.DataFrame({
    'date_from': ['2018-05-11', '2018-06-12'],
    'date_to': ['2018-06-12', '2018-06-25'],
    'days_from': [3, 2],
    'days_to': [12, 8],
    'koeff': [6.8714, 5.7825]
})

# 转换日期格式后调用函数
df['r_date'] = pd.to_datetime(df['r_date'])
res_df['date_from'] = pd.to_datetime(res_df['date_from'])
res_df['date_to'] = pd.to_datetime(res_df['date_to'])
result = test_lst(df, res_df)
print(result)

如果你的规则还需要额外匹配product等其他字段，直接在条件判断语句中添加对应判断逻辑即可。

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