如何在 pandas 中对多列向后应用滚动函数？

Question

给定这个简单的数据框：

df = pd.DataFrame(np.random.randint(0,100, size=(50, 4)), columns=list('ABCD'))

我正在尝试执行以下计算：

添加名为 B1、C2 和 D2 的三列，默认用 NaN 填充。
逐行检查A列接下来的5行，第一个大于20的行，然后B1、C2、D2列将填充该特定行的B、C、D列的内容.
如果 A 列接下来的 5 行中的 none 小于 20，则 B1、C2 和 D2 列将保持为 NaN。

我想到了这个方法：

def check_thresh(ser):
    dft = df.loc[ser.index]
    
    for _, row in dft.iterrows():
        if row['A'] > 20:
            return np.array([row['B'], row['C'], row['D']])
        
    return np.array([np.nan, np.nan, np.nan])

rol = df['A'].rolling(window=5)
df[['B1', 'C1', 'D1']] = rol.apply(check_thresh, raw=False)

但是，我面临以下问题：

它检查 previous 5 行，而不是 next 5 行。
性能很慢，而且我必须处理大型数据集。
它 returns 出现以下错误：TypeError: only size-1 arrays can be converted to Python scalars 将滚动功能应用于新列时。

我的方法有什么问题？你知道有更好的方法来处理这种情况吗？

Answer 1

我不确定这个实现是否经过优化或者是否正确，因为我没有完全理解这个问题并且没有预期输出的示例。

from numpy.lib.stride_tricks import sliding_window_view

WINDOWSIZE = 5
THRESHOLD = 20

# Equivalent to pd.rolling
m = sliding_window_view(df, (WINDOWSIZE, len(df.columns))).squeeze().astype(float)

# Extract 'A' column
A = m[:, :, 0]

# Get the first index whose value > THRESHOLD
argm = np.argmax(A > THRESHOLD, axis=1)

# True if all values <= THRESHOLD
amin = np.amin(A <= THRESHOLD, axis=1)

# Select rows in original array m
r = np.take_along_axis(m, argm[:, np.newaxis, np.newaxis], axis=1).squeeze()
r[amin] = np.nan

示例：

>>> df
    A   B   C   D
0   0   1   2   3
1   4   5   6   7
2   8   9  10  11
3  12  13  14  15
4  16  17  18  19
5  20  21  22  23
6  24  25  26  27
7  28  29  30  31
8  32  33  34  35
9  36  37  38  39

# df1 = pd.DataFrame(A).rename(columns='A{}'.format).assign(argm=argm, amin=amin)
# df2 = pd.DataFrame(r, columns=['A', 'B1', 'C1', 'D1'])

>>> pd.concat([df1, df2], axis='columns')
     A0    A1    A2    A3    A4  argm   amin     A    B1    C1    D1
0   0.0   4.0   8.0  12.0  16.0     0   True   NaN   NaN   NaN   NaN
1   4.0   8.0  12.0  16.0  20.0     0   True   NaN   NaN   NaN   NaN
2   8.0  12.0  16.0  20.0  24.0     4  False  24.0  25.0  26.0  27.0
3  12.0  16.0  20.0  24.0  28.0     3  False  24.0  25.0  26.0  27.0
4  16.0  20.0  24.0  28.0  32.0     2  False  24.0  25.0  26.0  27.0
5  20.0  24.0  28.0  32.0  36.0     1  False  24.0  25.0  26.0  27.0

注：最终dataframe的长度为len(df) - WINDOWSIZE + 1

如何在 pandas 中对多列向后应用滚动函数？

How to apply rolling function backwards with multiple columns in pandas?

python

numpy

dataframe

pandas

rolling-computation