Get rows before and after from an index in pandas dataframe-CodePudding

I want to get a specific amount of rows before and after a specific index. However, when I try to get the rows, and the range is greater than the number of indices, it does not return anything. Given this, I would like you to continue looking for rows, as I show below:

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

CodePudding user response：

You are mixing iloc and loc which is very dangerous. It works in your example because the index is sequentially numbered starting from zero so these two functions behave identically.

Anyhow, what you want is basically taking rows with wrap-around:

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)