Numpy中高效匹配二维数组行与另一数组行连续子数组的方法

这是个很常见的数组匹配需求，要做到高效的话，咱们得尽量避免Python层面的循环，利用NumPy的向量化操作和滑动窗口工具来搞定。下面直接上思路和代码：

核心思路 #

1. 先确定两个数组的列数：a的列数是我们要匹配的子部分长度，b的列数决定了每行能生成多少个连续子窗口。
2. 用滑动窗口工具把b的每一行拆成所有可能的、长度等于a列数的连续子数组——这一步不用复制数据，效率很高。
3. 利用广播机制，批量比较a的每一行和b的所有子窗口是否完全匹配。
4. 对每个a行和b行，只要存在至少一个匹配的子窗口，就记录这对索引。

代码实现 #

import numpy as np

# 示例输入数组
a = np.array([[ 25, 28], [ 84, 97], [105, 24], [ 28, 900]])
b = np.array([[ 25, 28, 84, 97], [ 22, 25, 28, 900], [ 11, 12, 105, 24]])

# 获取列数
m = a.shape[1]
n = b.shape[1]

# 生成b的滑动窗口视图（无数据复制，内存高效）
b_slide = np.lib.stride_tricks.sliding_window_view(b, window_shape=m, axis=1)

# 广播比较所有a行与b的子窗口，判断是否完全匹配
matches = (a[:, None, None, :] == b_slide[None, :, :, :]).all(axis=-1)

# 检查每个a行是否在某个b行中存在匹配的子部分
row_matches = matches.any(axis=-1)

# 提取所有匹配的(a行索引, b行索引)对
a_indices, b_indices = np.where(row_matches)
result = np.column_stack([a_indices, b_indices])

print(result)
# 输出符合预期：
# [[0 0]
#  [0 1]
#  [1 0]
#  [2 2]
#  [3 1]]

关键细节说明 #

• 滑动窗口的高效性：sliding_window_view是NumPy 1.20+新增的工具，它基于数组的 stride 机制生成视图，不会复制原始数据，所以内存占用极低，速度很快。如果你用的是旧版NumPy，可以用np.lib.stride_tricks.as_strided手动实现，但要注意计算正确的 stride 参数，避免越界。
• 向量化比较：通过广播把a和b_slide的维度对齐，一次性完成所有可能的匹配比较，比逐行循环快几个数量级，尤其是在数组规模较大的时候。
• 结果过滤：any(axis=-1)用来判断某个a行是否在某个b行中存在至少一个匹配的连续子部分，最后用np.where提取所有符合条件的索引对。

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