如何高效匹配NumPy浮点数组与整数数组的近似整数值？

针对大数据量的NumPy数组，核心是避免循环，利用向量化操作来提升效率，以下是两种高效实现方案：

方案一：使用np.isin快速判断存在性 #

np.isin是NumPy内置的向量化存在性检查函数，内部已做优化，适合大多数场景：

import numpy as np

x = np.array([0.43, 2.11, 3.67, 4.12, 7.05, 8.87, 9.1, 11, 11.67])
y = np.array([1, 3, 5, 9, 10])

# 确保y为整数类型（若原始y已是int可省略）
y = y.astype(int)

# 计算x中每个元素的向下/向上取整，转换为整数类型
x_floor = np.floor(x).astype(int)
x_ceil = np.ceil(x).astype(int)

# 分别判断floor和ceil结果是否在y中
floor_matched = np.isin(x_floor, y)
ceil_matched = np.isin(x_ceil, y)

# 最终结果：每个x元素是否满足floor或ceil在y中（布尔数组）
final_result = floor_matched | ceil_matched

# 若需要提取对应的匹配值（无匹配则返回np.nan）
matched_values = np.where(floor_matched, x_floor, np.where(ceil_matched, x_ceil, np.nan))

方案二：排序y后用np.searchsorted二分查找 #

如果y是静态数组（无需频繁修改），可以先排序，再用二分查找进一步提升性能（尤其当y规模极大时）：

import numpy as np

x = np.array([0.43, 2.11, 3.67, 4.12, 7.05, 8.87, 9.1, 11, 11.67])
y = np.array([1, 3, 5, 9, 10])

# 预处理：排序y
y_sorted = np.sort(y)

def check_in_sorted(arr, sorted_target):
    # 用二分查找获取插入位置
    idx = np.searchsorted(sorted_target, arr)
    # 检查位置对应的元素是否匹配，或前一个位置是否匹配（处理相等元素）
    return (idx < len(sorted_target)) & (sorted_target[idx] == arr) | (idx > 0) & (sorted_target[idx-1] == arr)

# 计算取整结果并检查
x_floor = np.floor(x).astype(int)
x_ceil = np.ceil(x).astype(int)

floor_matched = check_in_sorted(x_floor, y_sorted)
ceil_matched = check_in_sorted(x_ceil, y_sorted)

final_result = floor_matched | ceil_matched

效率说明 #

• 两种方案均为向量化操作，时间复杂度远低于循环遍历（循环为O(n*m)，向量化为O(n log m)）
• 若y频繁更新，优先用方案一；若y固定不变，方案二的排序成本可被多次查询摊薄，效率更高

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