如何修改numpy代码以识别一维数组中所有连续正值三元组的起始索引

原代码仅返回首个匹配的连续n个正值起始索引，未收集所有符合条件的位置，修改后代码如下：

import numpy as np

arr = np.array([1, -1, 5, 6, 3, -4, 2, 5, 9, 2, 1, -6, 8])

def find_consec_pos(arr, n):
    pos_mask = np.greater(arr, 0)
    # 使用valid模式卷积，每个结果对应一个完整长度为n的窗口的正值计数
    conv_res = np.convolve(pos_mask, np.ones(n, dtype=int), mode='valid')
    # 提取所有符合条件的起始索引转为列表
    valid_indices = np.where(conv_res >= n)[0].tolist()
    return valid_indices if valid_indices else None

print(find_consec_pos(arr, 3))

运行代码输出为[2, 6, 7, 8]，符合需求。

修改说明 #

• 卷积操作指定mode='valid'，过滤掉full模式下的边缘无效结果，返回的卷积结果索引直接对应滑动窗口的起始索引，无需额外偏移计算
• 用np.where()替代原逻辑的argmax()，收集所有符合条件的索引，而非仅取第一个匹配项
• 无符合条件结果时返回None，与原代码逻辑保持一致

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