Python 如何在pandas中使用多列向后滚动函数？

考虑到这个简单的数据帧：

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行中没有一行小于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)

然而，我面临以下问题：

它检查前5行，而不是后5行

性能相当慢，而且我必须处理大型数据集

它返回以下错误：

TypeError：将滚动函数应用于新列时，只能将size-1数组转换为Python标量

---

我的方法有什么问题？您知道更好的解决方案吗？

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

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

---

注意：最终数据帧的长度是

len（df）-WINDOWSIZE+1

这是一个多么难看的问题！然而，我认为5行仍然（仅仅）在使用硬编码的范围内。这样，您可以使用每个前瞻性行创建新列，并将其矢量化。换句话说，每一行也有列，告诉它下一行和下一行的值是什么。。。。forthAlso，我认为你目前的方法是O（N^2）。看起来您必须对每一行迭代整个df。顺便说一句，我不能运行您的示例。

TypeError:size-1数组可以转换为Python标量
df[['B1'，'C1'，'D1']]=rol.apply（check_thresh，raw=False）
@harrison4请编辑您的问题，并包含给定数据样本的预期输出。还有一件事，因为数据是使用随机样本生成的，所以在生成数据之前，您应该添加一个随机种子，以便我们能够重现您的数据帧。回答得很好！我真的很喜欢这个主意。
>>> 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