python - 使用 numpy 滚动最大值

Question

这将计算长度 K 的滑动窗口上 A 的“滚动最大值”(类似于滚动平均值):

import numpy as np
A = np.random.rand(100000)
K = 10
rollingmax = np.array([max(A[j:j+K]) for j in range(len(A)-K)])

但我认为它在性能方面远非最佳。

我知道 pandas 库有 rolling_max ，但在我的项目中，我不想使用这个新的依赖项。

问题:有没有一种简单的方法可以仅使用 numpy 计算滚动最大值？

Answer 1

我猜这个小技巧使用 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)