numpy如何不使用for循环查找A第0、2列匹配B的行索引

首先导入numpy依赖：

import numpy as np

再定义对应数组：

A = np.array([
    [1, 200, 15], [0, 600, 25], [2, 200, 20], [3, 100, 10],
    [1, 300, 5], [4, 400, 3], [0, 100, 12], [5, 300, 2],
    [2, 300, 25], [6, 100, 1], [1, 400, 74], [7, 300, 10],
    [2, 400, 15], [8, 200, 13], [3, 400, 29], [3, 600, 37],
    [0, 200, 20], [9, 300, 42], [5, 400, 30], [6, 200, 51],
    [7, 400, 21], [9, 500, 5], [10, 300, 9]
])
B = np.array([
    [0, 25.], [1, 74.], [2, 25.], [3, 37.], [4, 3.],
    [5, 30.], [6, 51.], [7, 21.], [8, 13.], [9, 42.], [10, 9.]
])

方法1：广播对比（中小数据量适用，代码最简洁） #

直接通过numpy向量化广播完成行匹配，无显式for循环：

# 提取A中需要对比的第0、2列，统一转为浮点型避免类型差异导致匹配失败
A_sub = A[:, [0, 2]].astype(float)
# 逐行判断A的子数组行是否存在于B的行列表中
match_mask = (A_sub[:, None] == B).all(axis=2).any(axis=1)
# 提取匹配的行索引
result = np.where(match_mask)[0]

输出结果为array([ 1, 5, 8, 10, 13, 15, 17, 18, 19, 20, 22], dtype=int64)，和预期完全一致。

方法2：结构化数组对比（大数据量适用，性能更优） #

如果A、B行数非常大，广播会产生高内存占用的中间矩阵，转一维结构化数组后对比效率更高：

A_sub = A[:, [0, 2]].astype(float)
# 定义统一的结构化数据类型
dtype = [('col0', float), ('col2', float)]
# 将二维数组转为一维结构化数组，每行对应一个结构化元素
A_view = A_sub.view(dtype).squeeze()
B_view = B.view(dtype).squeeze()
# 直接用一维匹配逻辑判断
match_mask = np.isin(A_view, B_view)
result = np.where(match_mask)[0]

两种方案的所有运算均为numpy内部实现的向量化操作，性能远高于手动遍历实现。

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