python - 从 pandas 数据框中的索引获取前后行

Question

我想获取特定索引之前和之后的特定行数。但是，当我尝试获取行并且范围大于索引数时，它不会返回任何内容。鉴于此，我希望您继续查找行，如下所示:

df = pd.DataFrame({'column': range(1, 6)})
    column
0   1
1   2
2   3
3   4
4   5

index = 2
df.iloc[idx]
3

# Now I want to get three values before and after that index.
# Something like this:
def get_before_after_rows(index):
    rows_before = df[(index-1): (index-1)-2]
    rows_after = df[(index+1): (index+1)-2]
    return rows_before, rows_after

rows_before, rows_after = get_before_after_rows(index)
rows_before
    column
0   1
1   2
4   5

rows_after
    column
0   1
3   4
4   5

Answer 1

您正在混合 iloc 和 loc 这是非常危险的。它在您的示例中有效，因为索引从零开始按顺序编号，因此这两个函数的行为相同。

无论如何，你想要的基本上是通过环绕来获取行:

def get_around(df: pd.DataFrame, index: int, n: int) -> (pd.DataFrame, pd.DataFrame):
    """Return n rows before and n rows after the specified positional index"""
    idx = index - np.arange(1, n+1)
    before = df.iloc[idx].sort_index()

    idx = (index + np.arange(1, n+1)) % len(df)
    after = df.iloc[idx].sort_index()

    return before, after

# Get 3 rows before and 3 rows after the *positional index* 2
before, after = get_around(df, 2, 3)