如何选取DataFrame中某列值距离特定值一定范围内的所有行？

Question

这里有一个示例数据框（DataFrame），我将用它来更好地说明我的问题：

import pandas as pd

df = pd.DataFrame(pd.np.random.rand(30, 3), columns=tuple('ABC'))
df['event'] = pd.np.nan
df.loc[10, 'event'] = 'ping'
df.loc[20, 'event'] = 'ping'
df.loc[19, 'event'] = 'pong'

我需要创建一个窗口，窗口的大小是 n 行，窗口的中心是每次出现的 ping。

换句话说，假设 i 是包含 ping 的那一行的索引。在每个 i 的情况下，我想选择 df.ix[i-n:i+n]。

因此，对于 n=3，我希望得到以下结果：

             A          B          C event
7    0.8295863  0.2162861  0.4856461   NaN
8     0.156646  0.4730667  0.9968878   NaN
9    0.6709413  0.4796197  0.8747416   NaN
10  0.09942329   0.154008  0.5761598  ping
11   0.7168143   0.678207  0.7281105   NaN
12   0.8915475  0.8013187  0.9049722   NaN
13   0.9545411  0.4844835  0.1645746   NaN
17   0.9909208  0.1091025  0.6582635   NaN
18   0.2536326  0.4324749  0.8001643   NaN
19   0.4734659  0.5582809  0.1221296  pong
20   0.7230407  0.6695843  0.3902591  ping
21   0.3624909  0.2685049  0.5484445   NaN
22  0.05626284  0.6113877  0.9131929   NaN
23   0.8312294  0.5694373  0.4325798   NaN

[14 rows x 4 columns]

有几点需要注意：

1. 我希望找到一个非迭代的解决方案。
2. 请注意，有一个 pong 的值，但我们不想以它为中心来创建窗口。不过，它在围绕第二个 ping 的结果中被包含了。

这该如何实现呢？

Answer 1

可能是：

>>> ts, n = df['event'] == 'ping', 3
>>> idx = ts.shift(n).fillna(False)  # +n rows
>>> for j in range(-n, n):  # -n to n-1 rows
...     idx |= ts.shift(j).fillna(False)
... 
>>> df[idx]

Answer 2

一种实现方法是使用嵌套的 np.where 语句。虽然代码看起来不太优雅，但能完成任务。

ping = pd.Series(np.where(df.event == 'ping', True,
                          np.where(df.event.shift(1) == 'ping', True,
                                   np.where(df.event.shift(-1) == 'ping', True, False))), index=df.index)

df[ping]

有没有人能帮我把 i=1 的情况推广到一般情况？

补充：其实它们不需要嵌套。这样就可以了：

ping = pd.Series(np.where((df.event == 'ping') | (df.event.shift(1) == 'ping') |
                      (df.event.shift(-1) == 'ping'), True, False), index=df.index)

Answer 3

In [17]: n = 3

选择一个索引范围，比如目标索引上下各加3（要注意不要超过数据框的最大或最小范围）。把这些索引连接在一起，然后去掉重复的。

In [18]: indexers = np.unique(np.concatenate([ np.arange(max(i-n,0),min(i+n,len(df))) for i in df[df.event=='ping'].index ]))

In [19]: indexers
Out[19]: array([ 7,  8,  9, 10, 11, 12, 17, 18, 19, 20, 21, 22])

选择它们。

In [20]: df.iloc[indexers]
Out[20]: 
             A           B          C event
7   0.03348742  0.05735324  0.1220022   NaN
8    0.9567363   0.6539097  0.8409577   NaN
9    0.3115902   0.4955503  0.1749197   NaN
10   0.6883777   0.6185107  0.7933182  ping
11   0.5185129   0.6533616  0.1569159   NaN
12   0.1196976   0.9638604  0.7318006   NaN
17  0.02897615   0.1224485  0.5706852   NaN
18  0.02409971   0.4715463  0.4587161   NaN
19   0.9070592   0.3371241  0.9543977  pong
20   0.8533369   0.7549413  0.5334882  ping
21   0.9546738   0.8203931  0.8543028   NaN
22  0.05691086   0.2402766  0.3922318   NaN

注意，你可能需要先执行 df.reset_index()（在选择之前，这样可以得到实际的行索引位置，而不是一个值）。

这里有个小问题，就是设置 'event' 列时会把所有内容都变成对象类型，具体可以查看 这里。你可以通过使用 df.convert_objects() 来缓解这个问题。