如何在NumPy数组中统计连续正、负、零元素？现有函数改造方案

我来帮你改造这个consecutive_counts函数，让它能同时统计连续负、零、正元素的数量。先说说原函数的局限：它仅将元素划分为“正”和“非正”两类，没法区分零和负数，所以我们首先要调整分类逻辑，把元素明确分成负、零、正三个类别。

改造后的函数代码 #

import numpy as np

def consecutive_counts(arr):
    '''
    Returns number of consecutive negative, zero, and positive numbers
    arr = np.array
    negative_counts, zero_counts, positive_counts = consecutive_counts(arr)
    '''
    # 给每个元素打标签：-1=负，0=零，1=正
    labels = np.where(arr < 0, -1, np.where(arr == 0, 0, 1))
    
    # 找到标签发生变化的位置（相邻元素标签不同的索引）
    idx = np.flatnonzero(labels[1:] != labels[:-1])
    
    # 计算每一段连续相同标签元素的长度
    if len(idx) == 0:
        # 整个数组都是同一种类型的元素
        count = np.array([arr.size])
    else:
        count = np.concatenate(([idx[0]+1], idx[1:] - idx[:-1], [arr.size - 1 - idx[-1]]))
    
    # 根据数组第一个元素的标签，对连续长度进行分组
    first_label = labels[0]
    negative = []
    zero = []
    positive = []
    
    # 遍历每一段的长度和对应的标签
    current_label = first_label
    for i, cnt in enumerate(count):
        if current_label == -1:
            negative.append(cnt)
        elif current_label == 0:
            zero.append(cnt)
        else:
            positive.append(cnt)
        # 更新下一段的标签（切换到分界点后的元素标签）
        if i < len(count)-1:
            current_label = labels[idx[i]+1]
    
    # 转换为NumPy数组返回
    return (np.array(negative), np.array(zero), np.array(positive))

函数说明 #

1. 标签分类：用嵌套的np.where给每个元素打上明确的类别标签，彻底区分负、零、正三类元素；
2. 找分界点：通过np.flatnonzero定位相邻元素类别变化的位置，这些位置就是连续段的分割点；
3. 计算连续长度：基于分界点计算每一段连续元素的长度，同时处理了数组全为同类型元素的边界情况；
4. 分组统计：从数组第一个元素的类别开始，依次把每段长度分配到对应类别的列表中，最后转为NumPy数组返回。

使用示例 #

针对你提供的输入数据（开头为0，后续是连续负数）：

# 转换为NumPy数组作为输入
arr = n.temp.p['50000'].values
# 获取三类连续元素的统计结果
neg_counts, zero_counts, pos_counts = consecutive_counts(arr)

比如开头的0.00会被统计到zero_counts中（长度为1），后续的连续负数会被收集到neg_counts里，如果数组中没有正数，pos_counts会返回空的NumPy数组。

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