Pandas中匹配随机SIM值对应累积概率列区间生成分类列的方法

现有一个pandas DataFrame，包含多列累积概率数据：A列数值介于0到1之间，从A到E列数值逐步递增，E列固定为1，数据样例如下：

ID          A    B    C    D    E     SIM
1:          0.49 0.64 0.86 0.97 1.00  0.98
2:          0.76 0.84 0.98 0.99 1.00  0.87
3:          0.32 0.56 0.72 0.92 1.00  0.12

其中SIM列为随机均匀分布数值。需要新增一列SIM_CAT，取值为SIM值所在区间右边界对应的列名，即当SIM值落在某两列的数值区间内时取右边界列的列名，预期结果样例如下：

ID          A    B    C    D    E     SIM  SIM_CAT
1:          0.49 0.64 0.86 0.97 1.00  0.98 E
2:          0.76 0.84 0.98 0.99 1.00  0.87 C
3:          0.32 0.56 0.72 0.92 1.00  0.12 A

由于A到E列逐行严格递增，我们可以直接用二分查找逻辑实现需求，有两种常用写法：

写法1：简洁易读版（适合中小数据量） #

import numpy as np

# 按顺序定义区间右边界对应的列名
cat_cols = ['A', 'B', 'C', 'D', 'E']

df['SIM_CAT'] = df.apply(
    lambda x: cat_cols[np.searchsorted(x[cat_cols], x['SIM'])],
    axis=1
)

逻辑说明：np.searchsorted会返回将SIM值插入有序序列后保持序列有序的索引位置，刚好对应第一个大于等于SIM值的右边界列的索引，直接取对应列名即可。

写法2：向量化高性能版（适合百万级以上大数据量） #

全程无循环，性能比apply写法高10倍以上：

import numpy as np

cat_cols = ['A', 'B', 'C', 'D', 'E']
cat_values = df[cat_cols].values
sim_values = df['SIM'].values.reshape(-1, 1)

# 批量找到每行第一个大于等于SIM值的位置
right_idx = np.argmax(cat_values >= sim_values, axis=1)
df['SIM_CAT'] = np.array(cat_cols)[right_idx]

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