您可以使用numpy将整个内容向量化。我已经创建了这个随机数据集,其中包含(约)12,000个介于0到199,999,999之间的索引和一个同样长的介于0到1之间的随机浮点数列表:
indices = np.unique(np.random.randint(2e8,size=(12000,)))
values = np.random.rand(len(indices))
然后,我构建了一个大小为 2*win+1
的索引数组,围绕每个 indices
,以及一个相应的数组,记录该点对移动平均值的贡献量:
win = 10
avg_idx = np.arange(-win, win+1) + indices[:, None]
avg_val = np.tile(values[:, None]/(2*win+1), (1, 2*win+1))
所有需要处理的就是重复的索引并将贡献添加到移动平均值中:
unique_idx, _ = np.unique(avg_idx, return_inverse=True)
mov_avg = np.bincount(_, weights=avg_val.ravel())
您现在可以获取移动平均值超过0.5的索引列表,例如:
unique_idx[mov_avg > 0.5]
就性能而言,首先将上述代码转换为函数:
def sparse_mov_avg(idx, val, win):
avg_idx = np.arange(-win, win+1) + idx[:, None]
avg_val = np.tile(val[:, None]/(2*win+1), (1, 2*win+1))
unique_idx, _ = np.unique(avg_idx, return_inverse=True)
mov_avg = np.bincount(_, weights=avg_val.ravel())
return unique_idx, mov_avg
以下是针对测试数据描述的几种窗口大小的时间表:
In [2]: %timeit sparse_mov_avg(indices, values, 10)
10 loops, best of 3: 33.7 ms per loop
In [3]: %timeit sparse_mov_avg(indices, values, 100)
1 loops, best of 3: 378 ms per loop
In [4]: %timeit sparse_mov_avg(indices, values, 1000)
1 loops, best of 3: 4.33 s per loop