基于Numpy的快速向量化方法：判断数组行是否匹配指定点

问题描述 #

基础场景 #

给定行/列坐标数组：

import numpy as np

idx_test_r = np.array([0, 0, 2, 0, 2, 4])
idx_test_c = np.array([0, 1, 0, 2, 2, 6])

# 合并为坐标数组，形状(6,2)
idx_test = np.stack((idx_test_r, idx_test_c), axis=-1)

以及两个目标点(0,0)、(2,2)（对应point_r = np.array([0,2])、point_c = np.array([0,2])），需要快速判断idx_test的每一行是否等于其中一个目标点，期望输出[True, False, False, False, True, False]，要求避免循环，保证极高的执行效率。

扩展场景 #

当idx_test_r和idx_test_c为二维数组时（例如idx_test_r = np.array([[0, 0, 2, 0, 2, 4], [1, 1, 3, 1, 3, 5]])，idx_test_c结构相同），现有实现速度偏慢，需要更优的向量化方案。

解决方案 #

一维坐标数组的极速实现 #

直接利用numpy的向量化操作，精准匹配目标点，完全避免循环：

import numpy as np

idx_test_r = np.array([0, 0, 2, 0, 2, 4])
idx_test_c = np.array([0, 1, 0, 2, 2, 6])
point_r = np.array([0, 2])
point_c = np.array([0, 2])

# 精准匹配两个目标点，逻辑或合并结果
result = ((idx_test_r == point_r[0]) & (idx_test_c == point_c[0])) | \
         ((idx_test_r == point_r[1]) & (idx_test_c == point_c[1]))

该方案直接对行列做元素级匹配，运算完全由numpy底层优化，速度远超循环实现。

二维坐标数组的高效优化 #

针对二维输入，推荐两种优化方案，可根据数据规模选择：

方案1：广播式全匹配

利用numpy的广播特性，一次性完成所有匹配判断，代码简洁高效：

import numpy as np

# 二维输入示例
idx_test_r = np.array([[0, 0, 2, 0, 2, 4], [1, 1, 3, 1, 3, 5]])
idx_test_c = np.array([[0, 1, 0, 2, 2, 6], [1, 2, 1, 3, 3, 6]])
point_r = np.array([0, 2])
point_c = np.array([0, 2])

# 合并目标点为(2,2)的数组
target_points = np.stack((point_r, point_c), axis=-1)
# 合并测试坐标为(2,6,2)的数组
test_coords = np.stack((idx_test_r, idx_test_c), axis=-1)

# 广播匹配：测试坐标与每个目标点逐一比对
matches = np.all(test_coords[:, :, None, :] == target_points[None, None, :, :], axis=-1)
# 对每个测试点，判断是否匹配任意目标点
result = np.any(matches, axis=-1)

方案2：候选过滤式匹配

当目标点数量较少时，先通过isin过滤掉不可能匹配的位置，再做精准匹配，可大幅减少运算量：

import numpy as np

# 二维输入示例
idx_test_r = np.array([[0, 0, 2, 0, 2, 4], [1, 1, 3, 1, 3, 5]])
idx_test_c = np.array([[0, 1, 0, 2, 2, 6], [1, 2, 1, 3, 3, 6]])
point_r = np.array([0, 2])
point_c = np.array([0, 2])
target_points = np.stack((point_r, point_c), axis=-1)

# 先筛选出行列都在目标点范围内的候选位置
candidate_mask = np.isin(idx_test_r, point_r) & np.isin(idx_test_c, point_c)
# 初始化结果数组
result = np.zeros_like(candidate_mask, dtype=bool)

# 仅对候选位置做精准匹配
for r, c in target_points:
    result |= ((idx_test_r == r) & (idx_test_c == c))

该方案内存占用更低，在测试数据量大、目标点少的场景下速度优势明显。

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