如何在np.select()中使用单值自定义函数处理DataFrame列

问题背景 #

已有如下代码生成的DataFrame：

import pandas as pd
import numpy as np

data = ['05/21/2021','05/21/2022','05/21/2023']    
df = pd.DataFrame(data, columns=['register_date'])
df['register_year'] = pd.to_datetime(df['register_date']).dt.year
df['study_year'] = 2022

DataFrame结构：

需求：创建duration_start列，规则如下：

• 当register_year < study_year时，值为1
• 当register_year > study_year时，值为0
• 当register_year == study_year时，计算注册日期到当年年末的天数占全年的比例

已定义条件：

register_year_lt_study_year = df['register_year'] < df['study_year']
register_year_gt_study_year = df['register_year'] > df['study_year']

以及单值计算函数：

def proportion_to_year_end(date):
    start = pd.to_datetime(date)
    year_end = pd.to_datetime('12/31/' + str(start.year))
    return (year_end - start).days/365

需要解决：如何在np.select()的默认值位置，高效生成与DataFrame同长度的比例列，避免繁琐的临时列或循环。

解决方案 #

方法1：直接在np.select中嵌入Series.apply() #

无需生成临时列，直接将df['register_date'].apply(proportion_to_year_end)作为np.select的默认参数传入，直接得到对应长度的列：

df['duration_start'] = np.select(
    [register_year_lt_study_year, register_year_gt_study_year],
    [1, 0],
    default=df['register_date'].apply(proportion_to_year_end)
)

执行后df结果：

方法2：将单值函数改为矢量化实现（更高效） #

避免apply()的逐行处理，用pandas矢量化日期操作提升性能：

def vectorized_proportion_to_year_end(dates):
    start_dates = pd.to_datetime(dates)
    year_ends = pd.to_datetime(start_dates.dt.year.astype(str) + '-12-31')
    return (year_ends - start_dates).dt.days / 365

# 直接传入整列使用
df['duration_start'] = np.select(
    [register_year_lt_study_year, register_year_gt_study_year],
    [1, 0],
    default=vectorized_proportion_to_year_end(df['register_date'])
)

这种方式利用pandas原生矢量化API，数据量越大，性能优势越明显。

方法3：用where条件分步赋值（更直观） #

不用np.select()，通过布尔索引分步赋值，逻辑清晰可读性强：

# 初始化列为比例值
df['duration_start'] = vectorized_proportion_to_year_end(df['register_date'])
# 覆盖小于的情况
df.loc[register_year_lt_study_year, 'duration_start'] = 1
# 覆盖大于的情况
df.loc[register_year_gt_study_year, 'duration_start'] = 0

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