基于Numpy的可变窗口大小滑动窗口均值计算优化求助

我生成了带有递增时间戳（timestamp）的数据，希望基于时间戳将数组划分为窗口，计算每个窗口内数据的均值，最终得到包含新"时间戳"与对应均值的新数组。当前代码可正常运行，但我认为存在更贴合Numpy风格的实现方式，希望移除while循环，利用数组已排序的特性通过np.where实现批量处理。

原代码如下：

# generating test data, first row timestamps, always counting up and random data 
import numpy as np
data = np.array([np.cumsum(np.random.randint(100, size=20)), np.random.randint(1, 5, size=20)])
print(data)

window_size = 200
overlap = 100

i, l_lim, u_lim = 0, 0, window_size
timestamps = []
window_mean = []
while u_lim < data[0, -1]:
    window_mean.append(np.mean(data[1, np.where((data[0, :] > l_lim) & (data[0, :] <= u_lim))]))
    timestamps.append(i)
    l_lim = u_lim - overlap
    u_lim = l_lim + window_size
    i += 1

print(np.array([timestamps, window_mean]))

可以利用Numpy的searchsorted快速定位窗口边界，结合add.reduceat批量计算均值，完全消除循环：

import numpy as np

# 生成测试数据
data = np.array([np.cumsum(np.random.randint(100, size=20)), np.random.randint(1, 5, size=20)])
window_size = 200
overlap = 100

timestamps_data = data[0]
values_data = data[1]

# 批量生成所有窗口的上下限
step = window_size - overlap
upper_lims = np.arange(window_size, timestamps_data[-1], step)
lower_lims = upper_lims - window_size

# 用searchsorted快速定位每个窗口的左右数据索引
left_idxs = np.searchsorted(timestamps_data, lower_lims, side='right')
right_idxs = np.searchsorted(timestamps_data, upper_lims, side='right')

# 构造分割点，用reduceat批量计算窗口总和
split_points = np.concatenate([left_idxs, [right_idxs[-1]]])
window_sums = np.add.reduceat(values_data, split_points)[:-1]

# 计算每个窗口的元素个数，再得到均值（处理空窗口情况）
window_counts = right_idxs - left_idxs
window_means = np.where(window_counts > 0, window_sums / window_counts, np.nan)

# 生成新的时间戳索引
new_timestamps = np.arange(len(upper_lims))

# 最终结果数组
result = np.array([new_timestamps, window_means])
print(result)

关键逻辑说明 #

• 批量生成窗口边界：通过np.arange一次性算出所有窗口的上下限，替代循环中的逐次更新
• 快速定位数据范围：searchsorted利用时间戳已排序的特性，以O(log n)效率定位每个窗口的左右索引，比循环内重复调用np.where更高效
• 批量计算均值：add.reduceat一次性完成所有窗口的求和操作，结合窗口元素个数得到均值，彻底避免循环迭代

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