NumPy中如何基于两个对齐数组高效实现查询数组的映射取值？

现有三个NumPy数组：x、y、q，其中x和y长度相同、元素按索引一一对应，q为查询数组。已知x和q内所有值唯一，需求如下：

• 遍历q中每个值，查找其在x中对应的索引位置
• 提取y中相同索引位置的值作为输出
• 若q中的值不存在于x中，对应位置返回np.nan

测试用例：

import numpy as np
x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
q = np.array([2, 0])

预期输出：array([ 5., nan])

原有双层循环实现时间复杂度为O(n*m)，数据量较大时运行效率极低，下面提供两种更高效的实现方案。

NumPy 原生实现（性能最优） #

简洁版（用到np.argwhere，适合中小数据量） #

直接通过广播做等值/近似匹配，用np.argwhere提取匹配位置的索引，逻辑直观代码短：

# 精确匹配场景把np.isclose换成==即可
match_mask = np.isclose(q[:, np.newaxis], x[np.newaxis, :])
match_pos = np.argwhere(match_mask)
out = np.full(len(q), np.nan, dtype=np.float64)
out[match_pos[:, 0]] = y[match_pos[:, 1]]

注意：该方法会生成形状为(len(q), len(x))的中间矩阵，若两个数组长度均超过10万，内存占用会明显升高，此时推荐用下面的高性能版本。

高性能版（适合大数据量，低内存占用） #

先对x排序，再用np.searchsorted做二分查找，时间复杂度为O(nlogn + mlogn)，内存占用极低：

# 对x的索引按x值排序
sorted_indices = np.argsort(x)
x_sorted = x[sorted_indices]
y_sorted = y[sorted_indices]
# 二分查找q中元素在排序后x中的位置
insert_pos = np.searchsorted(x_sorted, q)
# 处理边界防止越界
insert_pos[insert_pos == len(x_sorted)] = len(x_sorted) - 1
# 判断是否真的匹配：精确匹配换==，近似匹配用np.isclose
exist_mask = np.isclose(x_sorted[insert_pos], q)
out = np.full(len(q), np.nan, dtype=np.float64)
out[exist_mask] = y_sorted[insert_pos[exist_mask]]

Pandas 实现（代码最简洁） #

直接利用Pandas的索引映射能力，核心代码仅需两行，不需要手动处理索引逻辑，适合快速开发：

import pandas as pd
# 构建值到映射结果的索引序列
lookup_series = pd.Series(y, index=x)
# 按q重索引，不存在的值自动填充NaN
out = lookup_series.reindex(q).to_numpy()

如果需要处理浮点近似匹配，可以换用pd.merge_asof实现公差范围内的匹配，精确匹配场景下上述代码足够用。

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