高效向量化检测连续相同数值段的起始索引与频次（支持 NaN）-Python教程-PHP中文网

高效向量化检测连续相同数值段的起始索引与频次（支持 nan）

本文介绍一种完全向量化、无显式循环的 NumPy 方法，用于快速识别一维数组中连续相等元素段的起始位置及长度；该方法天然兼容 NaN（利用 np.nan != np.nan 特性），适用于热浪识别、事件序列统计等场景，性能较原始 while 循环提升数十倍。

本文介绍一种完全向量化、无显式循环的 NumPy 方法，用于快速识别一维数组中连续相等元素段的起始位置及长度；该方法天然兼容 NaN（利用 `np.nan != np.nan` 特性），适用于热浪识别、事件序列统计等场景，性能较原始 while 循环提升数十倍。

在科学计算与时间序列分析中，常需统计「连续事件」的发生频次与位置——例如热浪天数（标记为 1）、干旱期（非 NaN）、或任意重复数值块（如 [1,2,2,3,3,3] 中的 2,2 和 3,3,3）。原始实现采用嵌套 while 循环，逻辑清晰但严重依赖 Python 层迭代，在处理大规模 3D 数据（如 (1000, 500, 500) 气象网格）时成为显著瓶颈。

幸运的是，NumPy 提供了高效的向量化原语，可将该任务转化为边界检测 → 长度推导 → 条件过滤三步流水线，全程避免 Python 循环与条件分支：

✅ 核心向量化步骤详解

定位组边界（Group Boundaries）
```
boundaries = np.where(np.diff(aa) != 0)[0] + 1
```
- np.diff(aa) 计算相邻差值，对 NaN 自动返回 NaN（因 np.nan != np.nan 恒为 True，故 diff([1,1,np.nan,np.nan,1]) → [0, nan, 0, nan]）；
- != 0 生成布尔掩码，np.where(...)[0] 提取所有 True 的索引（即值变化处的前一个位置）；
- +1 将其校正为新组的起始索引（如 [1,1,1,NaN,NaN,1] → boundaries = [3, 5]）。
计算各组长度（Group Lengths）
```
positions = np.r_[0, boundaries, len(aa)]  # [0, 3, 5, 8]
lengths = np.diff(positions)                  # [3, 2, 3]
```
- np.r_[...] 构造包含「数组起点、所有边界、数组终点」的有序位置序列；
- np.diff 直接给出每段跨度（即连续相同值的个数）。
提取有效组的起始索引与长度（过滤单元素组）
```
idx = np.concatenate(([0], boundaries))[lengths >= 2]
freqs = lengths[lengths >= 2]
```
- 注意：首个组起始索引恒为 0，因此将 boundaries 前置 0 后切片对齐；
- lengths >= 2 掩码同步作用于索引与长度，仅保留长度 ≥2 的事件段（符合原始函数行为）。

? 完整可运行示例

import numpy as np

def get_frequency_of_events_vectorized(aa):
    if len(aa) == 0:
        return np.array([]), np.array([])
    if len(aa) == 1:
        return np.array([0]), np.array([1])

    # Step 1: Find group boundaries
    boundaries = np.where(np.diff(aa) != 0)[0] + 1

    # Step 2: Compute group lengths
    positions = np.r_[0, boundaries, len(aa)]
    lengths = np.diff(positions)

    # Step 3: Filter groups with length >= 2, extract start indices & frequencies
    mask = lengths >= 2
    start_indices = np.concatenate(([0], boundaries))[mask]
    frequencies = lengths[mask]

    return start_indices, frequencies

# 测试用例
aa1 = np.array([1,2,2,3,3,3,4,4,4,4,5,5,5,5,5])
idx1, freq1 = get_frequency_of_events_vectorized(aa1)
print("Input:", aa1)
print("Indices:", idx1)      # [ 1  3  6 10]
print("Frequencies:", freq1) # [2 3 4 5]

aa2 = np.array([1,1,1,np.nan,np.nan,1,1,np.nan])
idx2, freq2 = get_frequency_of_events_vectorized(aa2)
print("\nInput:", aa2)
print("Indices:", idx2)      # [0 5]
print("Frequencies:", freq2) # [3 2]

⚠️ 注意事项与最佳实践

NaN 处理是安全的：得益于 np.nan != np.nan 返回 True，np.diff() 在遇到 NaN 时自动触发边界判定，无需特殊预处理。
空/单元素数组保护：函数已添加边界检查，避免 np.diff() 在长度
不适用于“严格相等”自定义逻辑：若需基于容差比较（如浮点近似相等），应先用 np.round() 或 np.isclose() 预处理生成布尔掩码，再应用本方法。
扩展至高维数据：对 3D 数据（如 (time, lat, lon)），可沿指定轴（如 axis=0）使用 np.apply_along_axis，或改用 dask.array 实现分块并行化。