如何在numpy条件逻辑的default段返回列值实现pandas新列赋值

问题原因 #

np.select要求conditions条件列表和choices返回值列表的长度必须一一对应，你的代码里定义了4个判断条件，但只提供了2个返回值，长度不匹配导致逻辑错误。另外需要注意列名大小写要和实际DataFrame的列名保持一致，你的示例DataFrame列名为大写COLOR，原代码用了小写color也可能触发报错。

修正方案 #

方案1：对齐条件和返回值长度 #

直接补全返回值列表，和4个条件一一对应：

import pandas as pd
import numpy as np

condition = [
    (df['COLOR'] == 'grey'),
    (df['COLOR'] == 'white'),
    (df['COLOR'] == 'orange'),
    (df['COLOR'] == 'pink')
]
result = ['black', 'black', 'red', 'red']
df['FINAL_COLOR'] = np.select(condition, result, default=df['COLOR'])

方案2：合并同结果条件（更简洁） #

把返回相同结果的条件用isin合并，减少冗余代码：

import pandas as pd
import numpy as np

condition = [
    df['COLOR'].isin(['grey', 'white']),
    df['COLOR'].isin(['orange', 'pink'])
]
result = ['black', 'red']
df['FINAL_COLOR'] = np.select(condition, result, default=df['COLOR'])

两种方案的default参数直接传入原df['COLOR']即可实现不满足条件时保留原值的需求，运行后就能得到你预期的输出结果。

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