比较两个矩阵对应列元素，求匹配项在array_2对应列索引的向量化实现

实现思路 #

完全基于Numpy向量化操作实现，无Python层级显式for循环，大数组处理效率远高于循环方案：

• 利用Numpy广播机制扩展数组维度，实现逐列、逐元素的批量比较
• 单独兼容np.nan匹配逻辑（原生Numpy中np.nan == np.nan返回False，需单独判断）
• 用argmax提取首个匹配位置，配合掩码标记无匹配的场景

完整代码实现 #

依赖导入与示例输入 #

import numpy as np

# 示例输入数组
array_1 = np.array([[1, 2], [2, np.nan]])
array_2 = np.array([[2, np.nan], [1, 3]])

核心逻辑 #

# 扩展维度实现广播比较：shape从(M,N)、(K,N)变为(M,K,N)
comp = array_1[:, np.newaxis, :] == array_2[np.newaxis, :, :]
# 处理nan的匹配场景
nan_mask = np.isnan(array_1[:, np.newaxis, :]) & np.isnan(array_2[np.newaxis, :, :])
comp[nan_mask] = True
# 提取首个匹配位置，无匹配则标记为False
has_match = comp.any(axis=1)
res = np.where(has_match, np.argmax(comp, axis=1), False)

输出验证 #

print(res)
# 输出结果和预期一致：
# [[1 False]
#  [0 0]]

补充说明 #

• 若array_2某列存在多个重复匹配元素，默认返回第一个出现的匹配元素索引，符合常规查找逻辑
• 输出数组为object类型（同时包含整数和布尔值），若需要统一数据类型可根据业务场景自行调整输出规则

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