how do I calculate a rolling idxmax

前端 未结 5 1739
梦如初夏
梦如初夏 2020-11-30 12:16

consider the pd.Series s

import pandas as pd
import numpy as np

np.random.seed([3,1415])
s = pd.Series(np.random.randint(0, 10, 10         


        
5条回答
  •  佛祖请我去吃肉
    2020-11-30 12:42

    Here's an approach using broadcasting -

    maxidx = (s.values[np.arange(s.size-3+1)[:,None] + np.arange(3)]).argmax(1)
    out = s.index[maxidx+np.arange(maxidx.size)]
    

    This generates all the indices corresponding to the rolling windows, indexes into the extracted array version with those and thus gets the max indices for each window. For a more efficient indexing, we can use NumPy strides, like so -

    arr = s.values
    n = arr.strides[0]
    maxidx = np.lib.stride_tricks.as_strided(arr, \
                       shape=(s.size-3+1,3), strides=(n,n)).argmax(1)
    

提交回复
热议问题