Numpy：在含大量NaN的稀疏数组中查找指定序列（忽略NaN）

解决方案 #

利用数组的稀疏特性（大部分元素为NaN），先过滤出有效数据，再通过向量化操作快速匹配目标序列，彻底规避嵌套循环的高复杂度：

import numpy as np

def find_sequence_indices(sparse_array, sequence):
    # 过滤出非NaN元素及其原始索引
    non_nan_mask = ~np.isnan(sparse_array)
    valid_values = sparse_array[non_nan_mask].astype(int)
    valid_indices = np.arange(len(sparse_array))[non_nan_mask]
    
    seq_length = len(sequence)
    # 边界情况处理：序列为空或有效元素不足序列长度
    if seq_length == 0 or len(valid_values) < seq_length:
        return []
    
    # 生成滑动窗口并匹配目标序列
    sliding_windows = np.lib.stride_tricks.sliding_window_view(valid_values, seq_length)
    match_positions = np.where(np.all(sliding_windows == sequence, axis=1))[0]
    
    # 转换为原数组的起始/结束索引对
    return [(valid_indices[pos], valid_indices[pos + seq_length - 1]) for pos in match_positions]

思路解析 #

1. 预处理降维：用numpy向量化操作提取所有非NaN元素的值和它们在原数组中的索引，时间复杂度O(n)（n为原数组长度）。由于原数组以NaN为主，有效元素数量远小于n，后续处理的数据量会大幅压缩。
2. 滑动窗口匹配：通过sliding_window_view生成有效元素的滑动窗口（窗口长度等于目标序列长度），再用np.all向量化判断窗口是否与目标序列完全匹配，时间复杂度O(m)（m为有效元素数量）。
3. 索引映射：把匹配到的窗口位置转换为原数组的起始和结束索引，得到最终结果。

测试验证 #

运行题目中的示例：

sequence = np.array([1, 2, 3])
sparse_array1 = np.array([np.nan, 2, 3, np.nan, 1, 2, np.nan, 3, 2, np.nan, np.nan, 3, 1, np.nan, np.nan, np.nan, 2, np.nan, 3, 1, 1, np.nan])
sparse_array2 = np.array([3, 2, 1, 2, 1])

print(find_sequence_indices(sparse_array1, sequence))
# 输出: [(4, 7), (12, 18)]

print(find_sequence_indices(sparse_array2, sequence))
# 输出: []

结果完全符合预期。

复杂度对比 #

原三重循环方案的时间复杂度为O(n³)，本方案的时间复杂度为O(n + m)（m为有效元素数量，且m << n），在大型稀疏数组上的性能提升极为显著。

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