如何在Python的NumPy数组中搜索子向量？

很高兴帮你解决这个问题！你之前的代码能完美匹配完整行，但要找子向量的话，得换个思路——因为子向量长度和行长度不一致，没法直接用逐行全匹配的方式。下面给你两种实用的实现方案：

方法1：使用滑动窗口视图（推荐，简洁高效） #

从NumPy 1.20版本开始，官方提供了np.lib.stride_tricks.sliding_window_view函数，能快速生成每行的滑动窗口，让我们可以轻松对比子向量。

直接看代码：

import numpy as np

# 初始化你的原始数组
a = np.zeros([10,5])
a[0] = [5,6,4,8,5]
a[1] = [3,6,8,5,3]
a[2] = [3,2,1,5,3]
a[3] = [6,5,6,4,6]
a[4] = [3,4,7,6,3]
a[5] = [2,3,1,5,2]
a[6] = [1,1,3,2,1]
a[7] = [6,5,8,8,6]
a[8] = [5,4,9,7,5]
a[9] = [1,2,7,8,1]

search = [2,3,1,5]  # 目标子向量
search_arr = np.array(search)

# 生成每行的滑动窗口，窗口长度等于子向量长度
window_view = np.lib.stride_tricks.sliding_window_view(a, window_shape=len(search_arr), axis=1)

# 检查每个窗口是否和子向量完全匹配
matches = np.all(window_view == search_arr, axis=2)

# 获取所有匹配的行索引和子向量起始列索引
row_indices, col_start_indices = np.where(matches)

print(f"匹配的行索引：{row_indices}")
print(f"子向量在对应行的起始列位置：{col_start_indices}")

运行这段代码会输出：

匹配的行索引：[5]
子向量在对应行的起始列位置：[0]

正好对应第5行的前4个元素就是你要找的子向量。

原理很简单：sliding_window_view会把每行（长度5）转换成包含2个长度为4的窗口的结构（因为5-4+1=2），然后我们沿着窗口元素的轴做全匹配，最后用np.where提取所有匹配的位置。

方法2：手动构造滑动窗口（兼容旧版NumPy） #

如果你使用的NumPy版本低于1.20，没法用上面的函数，可以手动通过 stride 构造窗口，兼容性拉满：

import numpy as np

def find_subvector(arr, subvec):
    sub_len = len(subvec)
    arr_row_len = arr.shape[1]
    
    # 如果子向量比行还长，直接返回空结果
    if sub_len > arr_row_len:
        return np.array([]), np.array([])
    
    # 计算窗口的stride参数
    arr_strides = arr.strides
    window_strides = (arr_strides[0], arr_strides[1])
    window_shape = (arr.shape[0], arr_row_len - sub_len + 1, sub_len)
    
    # 用stride构造窗口视图
    window_view = np.lib.stride_tricks.as_strided(arr, shape=window_shape, strides=window_strides)
    
    # 匹配子向量并返回结果
    matches = np.all(window_view == subvec, axis=2)
    return np.where(matches)

# 调用函数查找子向量
row_idx, col_idx = find_subvector(a, search)
print(f"匹配行索引：{row_idx}, 子向量起始列：{col_idx}")

这个方法和第一种原理一致，只是手动计算了内存 stride，避免了额外的数据复制，效率同样很高。

额外提示 #

• 如果你的数组里有多个匹配的子向量位置，上面的方法会一次性找出所有结果；
• 记得先判断子向量长度是否大于行长度，避免出现无效计算。

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