如何用Numpy通用函数高效识别数组中的右向左阶梯结构？

问题描述 #

需要识别NumPy数组中是否存在从右向左的阶梯状结构：每一行的第一个有效1（相对于前一行的起始位置）必须位于前一行1的右侧，且从起始位置到该1之间不能有其他1（即起始位置到该1的前一个位置都是0）。现有循环实现，希望用NumPy通用函数优化。

原循环实现 #

import numpy as np

def pattern_found(sts):
    z = 0
    for i in range(sts.shape[0]): 
        x = np.argmax(sts[i, z:])
        z += x
        if not x: 
            return False
    return True

# 测试案例1：存在阶梯结构
states = np.array([[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                   [1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
                   [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0]])
print(pattern_found(states)) # 输出: True

# 测试案例2：不存在阶梯结构
states = np.array([[0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
                   [0, 1, 0, 0, 0, 0, 0, 1, 0, 1],
                   [0, 0, 0, 1, 1, 0, 0, 0, 0, 0]])
print(pattern_found(states)) # 输出: False

NumPy优化方案 #

由于该逻辑存在顺序依赖（每一行的起始位置由前一行的结果决定），完全无循环的向量化实现难度较高，但可以通过简化循环内的NumPy操作提升效率，同时保持逻辑清晰：

import numpy as np

def pattern_found_vectorized(sts):
    start = 0
    n_rows = sts.shape[0]
    
    for i in range(n_rows):
        # 提取当前行从start位置开始的子数组
        sub_row = sts[i, start:]
        # 找到子数组中第一个1的索引
        first_one_idx = np.argmax(sub_row)
        
        # 两种无效情况：
        # 1. 子数组全为0（无1可找）
        # 2. 第一个元素就是1（索引为0，不符合阶梯向右的要求）
        if first_one_idx == 0:
            if sub_row[0] != 1 or np.all(sub_row == 0):
                return False
            return False
        
        # 更新下一行的起始位置
        start += first_one_idx
    
    return True

# 测试案例验证
states1 = np.array([[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                   [1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
                   [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0]])
print(pattern_found_vectorized(states1)) # 输出: True

states2 = np.array([[0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
                   [0, 1, 0, 0, 0, 0, 0, 1, 0, 1],
                   [0, 0, 0, 1, 1, 0, 0, 0, 0, 0]])
print(pattern_found_vectorized(states2)) # 输出: False

优化说明 #

1. 利用NumPy的argmax快速定位子数组中第一个1的位置，比手动遍历高效得多
2. 合并无效条件判断，减少分支逻辑
3. 保持循环的状态跟踪（start变量），因为该逻辑本质是顺序依赖的，无法完全脱离循环实现

如果需要进一步提升性能，可以结合numba对循环进行JIT编译，但这已经超出纯NumPy通用函数的范畴。

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