如何用Numpy标准函数在(N,K)主数组中查找(M,K)搜索数组元素的索引？

问题描述 #

我有一个(N,K)维度的numpy.ndarray数组master，需要在其中查找多个K维数组的出现位置，这些K维数组共M个，存储在(M,K)维度的数组search中。

示例 #

当N=9、K=2、M=3时：

import numpy as np
master = np.array([[0, 0], [1, 0], [2, 0], [0, 1], [1, 1], [2, 1], [0, 2], [1, 2], [2, 2]])
search = np.array([[1, 2], [2, 0], [4, -2]])

期望得到结果如array([7, 2])或array([2, 7])，因为master数组的索引2和7对应的元素存在于search数组中。

尝试过的方法 #

• 使用np.isin的错误尝试：

  np.argwhere(np.all(np.isin(master, search), axis=1)).ravel()
  

  该方法返回的是所有值都来自search的索引，而非属于search中同一元素的索引，不符合需求。

• 效率极低的列表推导式：

  np.argwhere(np.any(np.array([[np.array_equal(master[i], search[j])
                                 for i in range(N)]
                                for j in range(M)]),
                     axis=0)).ravel()
  

  由于输入数据量较大，这种嵌套循环的方式速度过慢，无法满足性能要求。

解决方案 #

以下两种方法均仅使用Numpy标准函数实现，且性能远优于嵌套列表推导式：

方法一：广播匹配（适用于中小型数据） #

利用Numpy的广播机制实现逐行匹配：

# 扩展维度后进行广播式逐元素比较，再判断整行是否完全匹配
matches = np.all(master[:, np.newaxis, :] == search[np.newaxis, :, :], axis=2)
# 提取master中存在匹配项的索引
result = np.argwhere(np.any(matches, axis=1)).ravel()

原理：通过np.newaxis给master添加维度变为(N,1,K)，给search添加维度变为(1,M,K)，广播后生成(N,M,K)的逐元素比较结果；沿K轴判断整行是否完全相等，再沿M轴判断该行是否在search中存在匹配，最终提取对应索引。

方法二：结构化数组匹配（适用于大数据量） #

将多维数组转换为结构化数组，让np.isin直接匹配整行元素：

# 将每行K维数组打包为单个结构化元素
master_struct = master.view(np.dtype((np.void, master.dtype.itemsize * K)))
search_struct = search.view(np.dtype((np.void, search.dtype.itemsize * K)))
# 用np.isin匹配并提取索引
result = np.argwhere(np.isin(master_struct, search_struct)).ravel()

原理：把每个K维数组打包成字节型的“原子”元素，此时np.isin会将整行当作单个元素进行匹配，既避免了np.isin逐元素匹配的错误，又能利用Numpy的底层优化，性能更适合大数据量场景。

验证结果 #

对示例数据运行以上两种方法，都会得到符合预期的结果：

print(result)  # 输出 array([2, 7]) 或 array([7, 2])，顺序取决于匹配逻辑

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