Search code examples
pythonpandasdata-processing

Select a row based on column value and its previous 2 rows


+---+---+---+---+----+
| A | B | C | D | E  |
+---+---+---+---+----+
| 1 | 2 | 3 | 4 | VK |
| 1 | 4 | 6 | 9 | MD |
| 2 | 5 | 7 | 9 | V  |
| 2 | 3 | 5 | 8 | VK |
| 2 | 3 | 7 | 9 | V  |
| 1 | 1 | 1 | 1 | N  |
| 0 | 1 | 6 | 9 | V  |
| 1 | 2 | 5 | 7 | VK |
| 1 | 7 | 8 | 0 | MD |
| 1 | 5 | 7 | 9 | VK |
| 0 | 1 | 6 | 8 | V  |
+---+---+---+---+----+

i want to select a row based on column value and its two previous rows. For example in the given dataset (on the picture) I want to select row based on 'E' column value 'VK' and two previous rows of that selected row. So we should get a dataset like this:

+---+---+---+---+----+
| A | B | C | D | E  |
+---+---+---+---+----+
| 1 | 2 | 3 | 4 | VK |
| 1 | 4 | 6 | 9 | MD |
| 2 | 5 | 7 | 9 | V  |
| 2 | 3 | 5 | 8 | VK |
| 2 | 3 | 7 | 9 | V  |
| 1 | 1 | 1 | 1 | N  |
| 1 | 2 | 5 | 7 | VK |
| 1 | 7 | 8 | 0 | MD |
| 1 | 5 | 7 | 9 | VK |
+---+---+---+---+----+

Solution

  • 1st we need filter the dataframe until the last VK, then create the groupkey with cumsum , then do groupby head

    df=df.loc[:df.E.eq('VK').loc[lambda x : x].index.max()]
    df=df.iloc[::-1].groupby(df.E.eq('VK').iloc[::-1].cumsum()).head(3).sort_index()
    df
    Out[102]: 
       A  B  C  D   E
    0  1  2  3  4  VK
    1  1  4  6  9  MD
    2  2  5  7  9   V
    3  2  3  5  8  VK
    5  1  1  1  1   N
    6  0  1  6  9   V
    7  1  2  5  7  VK
    8  1  7  8  0  MD
    9  1  5  7  9  VK