pandas dataframe by boolean value, by index, and by integer

Question

我遇到了和这里类似的问题

我想要的是通过布尔索引（简单）获取 DataFrame 的一部分，并向后查看一些值，比如前一个索引，可能还有更多。不幸的是，在我可以获得实际整数位置之前，get_loc 链接问题中的建议答案使我的代码片段窒息（在以下片段中输入错误）。

以另一个问题的答案为例，这是我尝试过的：

df = pd.DataFrame(index=pd.date_range(start=dt.datetime(2015,1,1), end=dt.datetime(2015,2,1)), data={'a':np.arange(32)})
df.index.get_loc(df.index[df['a'] == 1])
*** TypeError: Cannot convert input to TimeStamp

之前的答案使用了 get_loc 的字符串，我只想传递一个普通索引值（这里是 DateTime）

Answer 1

使用切片:

import numpy as np
import pandas as pd
import datetime as DT
index = pd.date_range(start=DT.datetime(2015,1,1), end=DT.datetime(2015,2,1))
df = pd.DataFrame({'a':np.arange(len(index))}, index=index)

mask = df['a'] == 1
idx = np.flatnonzero(mask)[0]
lookback = 3
print(df.iloc[max(idx-lookback, 0):idx+1])

产量

             a
2015-01-08   7
2015-01-09   8
2015-01-10   9
2015-01-11  10

请注意，如果 idx-lookback 为负数，则索引引用 df 尾部附近的元素，就像 Python 列表一样：

In [163]: df.iloc[-3:2]
Out[163]: 
Empty DataFrame
Columns: [a]
Index: []

In [164]: df.iloc[0:2]
Out[164]: 
            a
2015-01-01  0
2015-01-02  1

因此，要获取相对于 df 头部的元素，请使用 max(idx-lookback, 0)。

使用布尔掩码:

如您所知，如果您有一个布尔数组或布尔系列，例如

mask = df['a'] == 10

你可以select对应的行

df.loc[mask]

如果您希望 select 之前或之后的行移动固定数量，您可以使用 mask.shift 移动掩码：

df.loc[mask.shift(-lookback).fillna(False)]

如果您希望 select lookback 前面的行，那么您可以通过将掩码与其移位合并来扩展掩码：

lookback = 3
for i in range(1, lookback):
    mask |= mask.shift(-i)

或者，等效地，使用 cumsum:

mask = (mask.shift(-lookback) - mask.shift(1)).cumsum().fillna(False).astype(bool)

for-loop 更清晰，但 cumsum 表达式更快，特别是如果 lookback 很大。

例如，

import numpy as np
import pandas as pd
import datetime as DT
df = pd.DataFrame(
    index=pd.date_range(start=DT.datetime(2015,1,1), end=DT.datetime(2015,2,1)), 
    data={'a':np.arange(32)})

mask = df['a'] == 10
lookback = 3
for i in range(1, lookback):
    mask |= mask.shift(-i)

# alternatively,
# mask = (mask.shift(-lookback) - mask.shift(1)).cumsum().fillna(False).astype(bool)

print(df.loc[mask])

产量

             a
2015-01-08   7
2015-01-09   8
2015-01-10   9
2015-01-11  10

pandas dataframe by boolean value, by index, and by integer

pandas dataframe by boolean value, by index, and by integer

python

pandas

indexing

boolean-expression