在数据框中选择下一个x行，不带iterrows

Question

假设我有一个如下所示的数据帧：

dftrial = pd.DataFrame({'col 1': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n'],
                        'col 2': [1, -1, -2, -3, 2, -1, 1, 5, 6, -2, -2, -2, -2, 1]})
dftrial

我想创建一个新的数据框，用于标识col 2为正的位置，然后将该行以及该行下方的下一个x行（在本例中，仅下一行），然后删除重复项

我可以使用iterrows来做这样的事情，如下所示

storage = []
for index, row in dftrial.iterrows():
    if row['col 2'] > 0:
        savethis = dftrial[index : index + 2]
        storage.append(savethis)
        
result = pd.concat(storage)
result.drop_duplicates(inplace = True)
display(result)

这将导致我想要的表格：

然而，我在一个有数十万行的数据帧上运行这段代码。按定义运行iterrows会遍历每一行，这会降低此过程的速度。有没有更快的方法来运行此代码？我知道有一些现有页面提供了iTerrow的替代方案，但它们似乎都没有解释我如何使用它来选择它下面的行

我得到的最接近的方法是使用lambda函数，如下所示，但该函数只接受col 2大于零的行，不包括以下行。我一直在努力使行的索引正确，任何帮助都将不胜感激

dftrial[dftrial.apply(lambda x: x['col 2'] > 0, axis = 1)]

Answer 1

要生成与描述不完全相同的示例，请执行以下操作：

>>> dftrial.loc[(dftrial['col 2'] > 0).rolling(2, min_periods=0).max() > 0]
   col 1  col 2
0      a      1
1      b     -1
4      e      2
5      f     -1
6      g      1
7      h      5
8      i      6
9      j     -2
13     n      1

如果我们严格按照您的描述（“下一列下一列两列”），那么：

>>> dftrial.loc[(dftrial['col 2'] > 0).rolling(
...     2, min_periods=0, closed='left').max() > 0]
   col 1  col 2
1      b     -1
2      c     -2
5      f     -1
6      g      1
7      h      5
8      i      6
9      j     -2
10     k     -2

解释

当然，关键是首先确定col 2为正的行：

>>> dftrial.assign(ispos=dftrial['col 2'] > 0)
   col 1  col 2  ispos
0      a      1   True
1      b     -1  False
2      c     -2  False
3      d     -3  False
4      e      2   True
5      f     -1  False
6      g      1   True
7      h      5   True
8      i      6   True
9      j     -2  False
10     k     -2  False
11     l     -2  False
12     m     -2  False
13     n      1   True

然后我们使用rolling().max()将该真理“出血”到下两行（不幸的是，Rolling没有.any()，否则我们将使用它而不是.max() > 0）：

>>> dftrial.assign(
...     ispos=dftrial['col 2'] > 0,
...     mask=(dftrial['col 2'] > 0).rolling(
...         2, min_periods=0, closed='left').max() > 0,
... )
   col 1  col 2  ispos   mask
0      a      1   True  False
1      b     -1  False   True
2      c     -2  False   True
3      d     -3  False  False
4      e      2   True  False
5      f     -1  False   True
6      g      1   True   True
7      h      5   True   True
8      i      6   True   True
9      j     -2  False   True
10     k     -2  False   True
11     l     -2  False  False
12     m     -2  False  False
13     n      1   True  False

然后，只需选择带有该遮罩的行