如何用Numpy原生函数加速值到连续区间的映射操作？

我有一个划分区间的levels数组：

#                  0    1    2    3   4
levels = np.array(( 0.2, 0.4, 0.6, 0.8 )) 

以及一个待映射的values数组，示例生成方式如下：

np.random.seed(20230204)
values = np.random.rand(5)

目前实现了一个执行速度较慢的映射函数：

def map_into_levels(values, levels):
    result = []
    for n in np.asarray(values):
        for r, level in enumerate(levels):
            if n <= level:
                break
        else:
            r += 1
        result.append(r)
    return result

该函数的作用是将values中的每个元素映射到levels划分的连续区间，示例运行结果如下：

In [153]: np.random.seed(20220204)
     ...: values = np.random.rand(6)
     ...: levels = np.array(( 0.2, 0.4, 0.6, 0.8 ))
     ...: result = map_into_levels(values, levels)
     ...: print(levels)
     ...: print(values)
     ...: print(result)
[0.2 0.4 0.6 0.8]
[0.00621839 0.23945242 0.87124946 0.56328486 0.5477085  0.88745812]
[0, 1, 4, 2, 2, 4]

请问可以使用哪个Numpy原生函数来加速这一操作？

解决方案 #

可以使用np.digitize()函数完全替代自定义的map_into_levels函数，它是Numpy原生的向量级运算函数，能大幅提升处理速度，且逻辑完全匹配需求。

使用示例 #

import numpy as np

np.random.seed(20220204)
values = np.random.rand(6)
levels = np.array((0.2, 0.4, 0.6, 0.8))

result = np.digitize(values, levels)
print(levels)
print(values)
print(result)

运行结果与自定义函数一致：

[0.2 0.4 0.6 0.8]
[0.00621839 0.23945242 0.87124946 0.56328486 0.5477085  0.88745812]
[0 1 4 2 2 4]

函数逻辑说明 #

np.digitize()的默认行为正好对应自定义函数的逻辑：

• 若元素x小于等于levels[0]，返回索引0；
• 若元素x落在levels[i-1] < x <= levels[i]区间，返回索引i；
• 若元素x大于所有levels中的值，返回len(levels)（即示例中的4）。

如果需要调整区间的闭合方向，可以通过right参数设置：

• 当right=True时，匹配规则变为levels[i-1] <= x < levels[i]，此时等于分界点的元素会被划分到右侧区间。

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