Select current row and the 3 rows above where a condition is met

Viewed 251

sample data:

df = pd.DataFrame({'user': ['Mike', 'Jim', 'Bob', 'Jane', 'Alice'], 
                   'income': [45000,55000, 40000, 50000, 42000],
                    'age' : [20,31,40,35,20]})

I want to be able to select 'user' and 'age' from the current row and the 2 rows above anyone who's name = Alice

here is what I have so far (not sure if any of it is right):

age2 = []
income2 = []

if df("user") in ['Alice']:
                age2.append(df.attrib.get("age"))
                income2.append(df.attrib.get("user"))

I want to have

age income
35   50000
40   40000
31   55000
4 Answers

IIUC groupby with cumsum, get first group and tail:

print (df.groupby(df["user"].eq("Alice").cumsum()).get_group(0).tail(3))

   user  income  age
1   Jim   55000   31
2   Bob   40000   40
3  Jane   50000   35

What you want is get a Series of the indexes that satisfy the condition you want. In your case the condition is user == Alice. I made extra entries in your dataframe to make it more clear. So:

import pandas as pd

df = pd.DataFrame({'user': ['Mike', 'Jim', 'Bob', 'Jane', 'Alice', 'Mike', 'Jim', 'Bob', 'Jane', 'Alice', 'Mike', 'Alice', 'Mike', 'Jim',], 
               'income': [45000,55000, 40000, 50000, 42000, 45000,55000, 40000, 50000, 42000, 55000, 40000, 50000, 42000],
                'age' : [20,31,40,35,20, 20,31,40,35,20, 40,35,20, 20]})

cond_idx = df[df.loc[:, "user"]=="Alice"].index

cond_idx now has the value

Int64Index([4, 9, 11], dtype='int64')

Now you can iterate through cond_idx and use .loc on dataframe to get the slices you want:

for idx in cond_idx:
    result = df.loc[idx-2:idx, ["age", "income"]]
    print(result)

Output:

   age  income
2   40   40000
3   35   50000
4   20   42000
   age  income
7   40   40000
8   35   50000
9   20   42000
    age  income
9    20   42000
10   40   55000
11   35   40000

Find the row indices that meet your condition

alices = df[df.user=='Alice'].index

Make an array for the previous three indices of each.

indices = alices.values[:,None] - [3,2,1]

Use those indices with .iloc

print(df.iloc[indices.ravel()])

Using Giorgos Livanos data

>>> indices
array([[ 1,  2,  3],
       [ 6,  7,  8],
       [ 8,  9, 10]], dtype=int64)
>>> indices.ravel()
array([ 1,  2,  3,  6,  7,  8,  8,  9, 10], dtype=int64)
>>>

pandas.Index.values
numpy Array Broadcasting

I'm not sure if this is what you are looking for, but I believe this is what you are attempting to do in your code:

df = pd.DataFrame({'user': ['Mike', 'Jim', 'Bob', 'Jane', 'Alice'], 
                   'income': [45000,55000, 40000, 50000, 42000],
                    'age' : [20,31,40,35,20]})

df2=[]

def search(name):
    for i in range(len(df['user'])):
        if name == df['user'][i]:
            df2=pd.DataFrame({'age':[df['age'][i]], 'income':[df['income'][i]]})
            print(df2)

search('Alice')
Related