在Python（Pandas / Numpy）中，如何使用条件和特定块大小对df进行子集化？

Question

我有一个df

A = pd.DataFrame([[1, 5, 2, 0], [2, 4, 4, 0], [3, 3, 1, 1], [4, 2, 2, 0], [5, 1, 4, 0], [2, 4, 4, 0], [3, 3, 1, 1], [4, 2, 2, 0], [5, 1, 4, 0]],
                 columns=['A', 'B', 'C', 'D'], index=[1, 2, 3, 4, 5, 6, 7, 8, 9])

我希望能够根据以下规则对数据帧进行子集化：选择列“D”值为1的行，并且还包括它们之上的两个行（块大小= 3）。

如果我在df示例中应用规则，则输出应为：

   A  B  C  D
1  1  5  2  0
2  2  4  4  0
3  3  3  1  1
5  5  1  4  0
6  2  4  4  0
7  3  3  1  1

由于

Answer 1

这适用于任何chunk尺寸：

>>> chunk, mask = 3, A['D'] == 1
>>> mask -= mask.shift(-chunk).fillna(0)
>>> A[mask[::-1].cumsum() > 0]
   A  B  C  D
1  1  5  2  0
2  2  4  4  0
3  3  3  1  1
5  5  1  4  0
6  2  4  4  0
7  3  3  1  1