在NumPy中寻找含最短非零序列的数组的最高效方法

给定的二维数组每行都是左侧连续1、右侧（可能没有）连续0的结构，要找非零序列最短的子数组，用NumPy的矢量化操作可以避免Python循环，直接依托C级处理提升效率，具体实现如下：

实现步骤 #

• 获取每行第一个0的索引：利用np.argmin，因为0是每行中的最小值，argmin会返回该行第一个最小值的位置；对于全1的行，argmin会返回0（所有元素都是1，第一个元素的索引为0）。
• 修正全1行的长度：全1行的非零序列长度等于数组的列数，用np.where把argmin返回0的行替换成列数z.shape[1]。
• 定位最短序列的子数组：计算修正后的长度数组的最小值，再找到该最小值对应的行索引即可。

完整代码示例 #

import numpy as np 

z = np.array(
    [
     [1, 1, 0, 0, 0, 0],
     [1, 1, 1, 1, 1, 0],
     [1, 1, 1, 0, 0, 0],
     [1, 1, 1, 1, 1, 1],
    ]
)

# 获取每行第一个0的索引
first_zero_idx = z.argmin(axis=1)
# 修正全1行的非零序列长度
non_zero_lengths = np.where(first_zero_idx == 0, z.shape[1], first_zero_idx)
# 找到最短长度对应的行索引
min_length = non_zero_lengths.min()
shortest_row_idx = np.where(non_zero_lengths == min_length)[0][0]
# 获取目标子数组
shortest_row = z[shortest_row_idx]

print(f"最短非零序列长度：{min_length}")
print(f"对应的子数组：{shortest_row}")

代码说明 #

• z.argmin(axis=1)：沿行方向查找最小值的索引，对于含0的行，这个索引就是非零序列的长度；全1行则返回0。
• np.where(...)：把全1行的长度修正为数组列数，因为这类行没有0，非零序列就是整行的长度。
• 最后通过min()找到最短长度，再用np.where定位到对应行索引，就能快速得到目标子数组。

这种方法全程用NumPy矢量化操作，没有Python循环，完全依托C级运算处理，在处理大规模数组时效率远高于逐行遍历的常规方法。

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