如何使用numpy向量化方法将浮点数组映射为np.unique对应排序整数索引

实现方法 #

你要的功能有两种常用的numpy向量化实现方案，根据你使用的唯一值数组来源选择即可：

方案1：唯一值数组从当前输入数组计算得到 #

直接调用np.unique时指定return_inverse=True参数，返回的第二个结果就是你需要的索引数组，完全符合你第一个示例的需求：

import numpy as np

# 第一个示例
A = np.array([0.1, 0.0, -0.2, -0.1, 0.1, 0.5])
_, B = np.unique(A, return_inverse=True)
print(B) # 输出 [3 2 0 1 3 4]

方案2：唯一值数组是预先定义的全量列表 #

如果你需要匹配的唯一值数组不是从当前输入数组计算得到（比如你补充示例里的7元素全量区间值），使用np.searchsorted实现映射即可：

# 补充示例
pre_defined_unique = np.array([-0.075, -0.05 , -0.025,  0.   ,  0.025,  0.05 ,  0.075])
C = np.array([0.025, 0.025, -0.025, -0.025, -0.025, 0.075, -0.025, -0.075, 0.075, -0.075])
D = np.searchsorted(pre_defined_unique, C)
print(D) # 输出 [4 4 2 2 2 6 2 0 6 0]

注意：浮点数存在精度误差问题，如果输入数组是计算得到的非精确值，建议先按照精度要求做四舍五入处理后再执行映射，避免匹配错误。

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