如何用NumPy向量化高效实现时间连续性检测函数？

需求说明 #

• 现有形状为(N, T, C)的3D NumPy数组（示例中N=3、T=5、C=1；实际场景N=100000、T=24、C=24）
• 数组中指定列存储表示小时的整数，要求每个band（即数组第一个维度的元素，形状(T,C)）满足：
  • 包含完整的0到T-1的所有整数（无重复、无缺失）
  • 整数序列是连续循环的正确顺序（比如[3,4,0,1,2]符合要求，[1,2,4,4,0]不符合）
• 需要全程用NumPy向量化操作筛选有效band，避免高内存开销

最小可复现代码 #

import numpy as np
# 定义数组形状：(band数量, 时间步长, 特征列数)
shape = (3, 5, 1)
# 创建随机数组
random_array = np.random.randint(0, 10, size=shape)
# 赋值符合要求的band
random_array[0, :, :] = [[0], [1], [2], [3], [4]]
random_array[1, :, :] = [[3], [4], [0], [1], [2]]
# 赋值不符合要求的band（存在重复和缺失）
random_array[2, :, :] = [[1], [2], [4], [4], [0]]

向量化解决方案 #

def filter_valid_bands(arr, time_col_idx=0, num_hours=None):
    # 获取时间步长T，未指定则默认用数组的时间维度长度
    T = arr.shape[1] if num_hours is None else num_hours
    # 提取所有band的时间列，转为(N, T)形状的2D数组
    time_sequences = arr[:, :, time_col_idx].reshape(-1, T)
    
    # 检查条件1：序列包含0到T-1的所有整数（无重复无缺失）
    sorted_seq = np.sort(time_sequences, axis=1)
    valid_full_set = np.all(sorted_seq == np.arange(T), axis=1)
    
    # 检查条件2：序列是连续循环的正确顺序
    # 计算相邻元素的差值，补充首尾循环的差值
    diffs = np.diff(time_sequences, axis=1)
    wrap_diff = time_sequences[:, 0] - time_sequences[:, -1]
    all_diffs = np.concatenate([diffs, wrap_diff[:, np.newaxis]], axis=1)
    # 合法差值只能是1（正常递增）或-(T-1)（循环跳转，比如T=5时4→0的差为-4）
    valid_order = np.all((all_diffs == 1) | (all_diffs == -(T-1)), axis=1)
    
    # 同时满足两个条件的band索引
    valid_indices = valid_full_set & valid_order
    
    # 返回筛选后的数组
    return arr[valid_indices]

# 测试函数
filtered_arr = filter_valid_bands(random_array, num_hours=5)
print(filtered_arr.shape)  # 输出(2,5,1)，符合预期

思路解释 #

1. 提取时间序列：从3D数组中取出目标时间列，转为(N,T)的2D数组，便于向量化操作
2. 完整性检查：对每个序列排序后，与0到T-1的标准序列对比，完全一致则说明无重复无缺失
3. 连续性检查：计算序列相邻元素（含循环首尾）的差值，所有差值符合1或-(T-1)则说明顺序正确
4. 筛选结果：通过逻辑与操作获取同时满足两个条件的band索引，直接筛选返回有效数组

这种方法全程用NumPy向量化实现，避免了循环和Pandas的内存开销，处理百万级band也能保持高效。

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