含0、1、2的大型二维数组：连续0筛选性能优化咨询

我有一个包含0、1、2的二维NumPy数组，列数极多。需要筛选出连续0的长度不超过指定数值的行。当前做法是将数组转为字符类型，合并每行的列后用正则表达式过滤，但类型转换和字符拼接步骤速度极慢，希望找到能提升一个数量级速度的方法，或者完全不同的解决方案。

当前使用的低效代码：

import re
import numpy as np

n=100
k = 1000
x = np.random.choice([0,1,2], replace=True, size=(n,k))
s = np.apply_along_axis(lambda t: ''.join(t) , 1, x.astype(str))

N_ramp=3
mask = [re.search(r'[12]0{1,'+str(N_ramp)+r'}[12]', i) is None for i in s]

直接基于NumPy的数值向量化操作实现，完全规避字符串转换和正则匹配的开销，速度能提升一个数量级以上：

import numpy as np

n = 100
k = 1000
x = np.random.choice([0,1,2], replace=True, size=(n,k))
N_ramp = 3

# 1. 标记所有非0元素（1/2视为非0，对应True；0对应False）
non_zero = x != 0

# 2. 计算每行中连续0的最大长度
# 在每行首尾添加True，统一处理行首/行尾的连续0情况
pad_non_zero = np.concatenate([np.ones((n,1), dtype=bool), non_zero, np.ones((n,1), dtype=bool)], axis=1)
# 获取每行中非0元素的列索引
indices = np.where(pad_non_zero)[1].reshape(n, -1)
# 相邻非0索引的差值减1，就是两段非0之间连续0的长度
zero_runs = indices[:, 1:] - indices[:, :-1] - 1
# 提取每行的最大连续0长度
max_zero_runs = zero_runs.max(axis=1)

# 3. 生成筛选掩码：保留连续0长度不超过N_ramp的行
mask = max_zero_runs <= N_ramp

方案优势 #

• 全程使用NumPy原生向量化运算，避免了Python循环和字符串处理的巨大开销
• 利用数组索引和差分计算连续区间，时间复杂度为O(n*k)，比正则匹配的字符串操作高效得多，尤其适合列数极多的场景

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