Pandas：如何解决np.select生成数组与DataFrame维度不匹配问题

原代码存在的错误 #

• 未正确生成DataFrame副本：data2 = data是对象引用而非独立副本，修改data2会同步修改原始data，需用.copy()方法生成独立副本
• 多余的groupby操作：qoe = data.groupby('world_rank')生成的是GroupBy对象，完全不符合新列的取值要求，后续将该对象插入DataFrame会直接触发长度不匹配错误
• np.select的条件构造错误：np.select要求每个条件是和DataFrame行数长度一致的布尔数组，你给每个条件加了.any()，会将整列布尔值聚合为单个布尔值，导致np.select返回长度为1的数组，和DataFrame行数不匹配；此外第二个条件中.any未加括号属于语法错误
• 重复添加同名列：先通过data2['qoe'] = ...添加了列，又调用insert再次插入同名qoe列，会触发列名重复冲突

修复后的可运行代码 #

import pandas as pd
import numpy as np
data = pd.read_csv("data/cwurData.csv")
# 生成独立副本
data2 = data.copy()
# 构造符合要求的条件数组，无需加any()
conditions = [
    data2['world_rank'] < 100,
    (data2['world_rank'] >= 100) & (data2['world_rank'] <= 300),
    data2['world_rank'] > 300
]
choices = ['High', 'Medium', 'Low']
# 直接生成qoe列数组
qoe_col = np.select(conditions, choices)
# 插入到第5位（索引从0开始，位置5就是第6列）
data2.insert(5, 'qoe', qoe_col)
# 校验结果
print(data2.axes)
data2.head(n=500)

注意事项 #

原逻辑中100、300两个边界值没有覆盖，上述代码将100到300的闭区间归为Medium，如果需求是左开右闭/左闭右开自行调整比较符号即可。

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