Полагаю, этот маленький трюк с использованием strides
и as_strided
сделает работу:
def max_rolling1(a, window,axis =1):
shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
strides = a.strides + (a.strides[-1],)
rolling = np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
return np.max(rolling,axis=axis)
и для целей совместного использования я определил другую функцию на основе вашего алгоритма:
def max_rolling2(A,K):
rollingmax = np.array([max(A[j:j+K]) for j in range(len(A)-K)])
return rollingmax
и сравнение по timeit
на моем ноутбуке:
с:
A = np.random.rand(100000)
K = 10
%timeit X = max_rolling2(A,K)
170 ms ± 19.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit X = max_rolling1(A,K)
> 3.75 ms ± 479 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)