您可以使用 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