如何基于多条件为DataFrame新增列？现有代码报错求优化

原代码的问题 #

1. 导入numpy时未设置别名np，但后续使用了np.select，会触发NameError；
2. np.select要求conditions和values的长度严格匹配，你这里10个条件对应1个值，会直接报错；
3. 未设置默认值，不满足条件的行会被赋值为nan，不符合你要赋值'0'的需求。

高效写法推荐 #

方法1：合并条件的向量化判断（性能最优） #

把同一年的月份合并成列表，用isin简化条件，再通过np.where直接生成目标列：

import numpy as np

# 合并所有满足条件的年月逻辑
lockdown_condition = ((final_df['Year'] == 2020) & final_df['Month'].isin([3,4,5,6,10,11,12])) | \
                     ((final_df['Year'] == 2021) & final_df['Month'].isin([1,2,3]))

# 生成Lockdown列，满足条件赋值'1'，否则'0'
final_df['Lockdown'] = np.where(lockdown_condition, '1', '0')

方法2：元组匹配（可读性强） #

将需要标记的年月组合成元组列表，通过向量化的元组匹配实现：

import numpy as np

# 定义所有封锁年月组合
lockdown_periods = [(2020,3), (2020,4), (2020,5), (2020,6),
                    (2020,10), (2020,11), (2020,12),
                    (2021,1), (2021,2), (2021,3)]

# 组合年月为元组并判断是否在列表内
final_df['Lockdown'] = np.where(
    final_df[['Year', 'Month']].apply(tuple, axis=1).isin(lockdown_periods),
    '1',
    '0'
)

方法3：日期区间判断（适合后续日期操作） #

把年月转换成日期格式，通过区间判断实现，适合需要后续做日期相关处理的场景：

import pandas as pd
import numpy as np

# 生成每月第一天的日期列（临时用）
final_df['YearMonth'] = pd.to_datetime(final_df[['Year', 'Month']].assign(Day=1))

# 定义两个封锁区间
period1 = (final_df['YearMonth'] >= '2020-03-01') & (final_df['YearMonth'] <= '2020-06-01')
period2 = (final_df['YearMonth'] >= '2020-10-01') & (final_df['YearMonth'] <= '2021-03-01')

# 生成Lockdown列
final_df['Lockdown'] = np.where(period1 | period2, '1', '0')

# 不需要临时列可删除
final_df.drop('YearMonth', axis=1, inplace=True)

性能说明 #

• 方法1的向量化判断速度最快，适合百万级以上的大数据量；
• 方法2代码直观，适合小数据量场景；
• 方法3灵活度高，适合需要扩展日期相关逻辑的场景。

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