仅使用迭代器的解决方案
foslok的解决方案绝对没问题,但我想尝试一下使用生成器的版本。它只在遍历原始列表时存储一个长度为window_size的双端队列,然后找到n_largest值并计算其平均值。
import itertools as it
from collections import deque
from heapq import nlargest
from statistics import mean
def windowed(iterable, n):
_iter = iter(iterable)
d = deque((it.islice(_iter, n)), maxlen=n)
yield tuple(d)
for i in _iter:
d.append(i)
yield tuple(d)
a = [3, 5, 2, 7, 5, 3, 6, 8, 4]
means = [mean(nlargest(2, w)) for w in windowed(a, 3)]
print(means)
结果:
[4, 6, 6, 6, 5.5, 7, 7]
因此,要更改元素数量(窗口大小)或n个最大元素,只需更改相应函数的参数即可。这种方法还避免了切片的使用,因此可以更轻松地应用于您无法或不想切片的可迭代对象。
时间记录
def deque_version(iterable, n, k):
means = (mean(nlargest(n, w)) for w in windowed(iterable, k))
for m in means:
pass
def tee_version(iterable, n, k):
means = (mean(nlargest(n, w)) for w in windowed_iterator(iterable, k))
for m in means:
pass
a = list(range(10**5))
n = 3
k = 2
print("n={} k={}".format(n, k))
print("Deque")
%timeit deque_version(a, n, k)
print("Tee")
%timeit tee_version(a, n, k)
n = 1000
k = 2
print("n={} k={}".format(n, k))
print("Deque")
%timeit deque_version(a, n, k)
print("Tee")
%timeit tee_version(a, n, k)
n = 50
k = 25
print("n={} k={}".format(n, k))
print("Deque")
%timeit deque_version(a, n, k)
print("Tee")
%timeit tee_version(a, n, k)
result:
n=3 k=2
Deque
1.28 s ± 3.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Tee
1.28 s ± 16.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n=1000 k=2
Deque
1.28 s ± 8.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Tee
1.27 s ± 2.92 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n=50 k=25
Deque
2.46 s ± 10.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Tee
2.47 s ± 2.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
显然,itertools的tee和deque并没有太大区别。