How to get a sequence of previous and next rows for each row in pandas

user_11666502 0
Viewed 77

I have a pandas df like this:

x     y   
1     0
5     1
3     0
2     0
5     1
6     0
1     0
4     0
3     1

I'm trying to create a new column that takes makes a list of the last 2, current, and next 2 elements in x and y. It should look like this:

x     y     seq
1     0     [(nan, nan), (nan, nan), (1, 0), (5,1), (3,0)]
5     1     [(nan, nan), (1, 0), (5, 1), (3,0), (2,0)]
3     0     [(1, 0), (5, 1), (3, 0), (2,0), (5,1)]
2     0     [(5, 1), (3, 0), (2, 0), (5,1), (6,0)]
5     1     [(3, 0), (2, 0), (5, 1), (6,0), (nan, nan)]
6     0     [(2, 0), (5, 1), (6, 0), (nan, nan),(nan, nan)]

I wrote this:

def sequences(df):

    back2 = (df.x.shift(2), df.y.shift(2))
    back1 = (df.x.shift(1), df.y.shift(1))
    current = (df.x, df.y)
    forward1 = (df.x.shift(-1), df.y.shift(-1))
    forward2 = (df.x.shift(-2), df.y.shift(-2))

    return [back2, back1, current, forward1, forward2]

df['data_sequence'] = df.apply(sequences, axis=1)

but .shift() fails in df.apply() because it treats each item like an int, not an element in a series. How do I accomplish this?

pandas python
2 Answers

You could do:

xy = list(zip(df['x'], df['y']))
xy = [(np.nan, np.nan)]*2 + xy + [(np.nan, np.nan)]*2
df['seq'] = [xy[i:i+5] for i in range(len(df))]

df:

    x   y
0   1   0
1   5   1
2   3   0
3   2   0
4   5   1
5   6   0

Output:

    x   y                                                seq
0   1   0   [(nan, nan), (nan, nan), (1, 0), (5, 1), (3, 0)]
1   5   1       [(nan, nan), (1, 0), (5, 1), (3, 0), (2, 0)]
2   3   0           [(1, 0), (5, 1), (3, 0), (2, 0), (5, 1)]
3   2   0           [(5, 1), (3, 0), (2, 0), (5, 1), (6, 0)]
4   5   1       [(3, 0), (2, 0), (5, 1), (6, 0), (nan, nan)]
5   6   0   [(2, 0), (5, 1), (6, 0), (nan, nan), (nan, nan)]
user_14103030

Following the approach that you are doing, less clean than SomeDude's code, but does the work:

def get_sequence(row, df):
    idx = row.name
    output = []
    for i in range(-2, 3):
        if 0 <= idx+i < df.shape[0]:
            output.append((df.iloc[idx+i].x, df.iloc[idx+i].y))
        else:
            output.append((np.nan, np.nan))
    return output

df["sequence"] = df.apply(lambda row: get_sequence(row, df), axis=1)

0    [(nan, nan), (nan, nan), (1, 0), (5, 1), (3, 0)]
1        [(nan, nan), (1, 0), (5, 1), (3, 0), (2, 0)]
2            [(1, 0), (5, 1), (3, 0), (2, 0), (5, 1)]
3            [(5, 1), (3, 0), (2, 0), (5, 1), (6, 0)]
4            [(3, 0), (2, 0), (5, 1), (6, 0), (1, 0)]
5            [(2, 0), (5, 1), (6, 0), (1, 0), (4, 0)]
6            [(5, 1), (6, 0), (1, 0), (4, 0), (3, 1)]
7        [(6, 0), (1, 0), (4, 0), (3, 1), (nan, nan)]
8    [(1, 0), (4, 0), (3, 1), (nan, nan), (nan, nan)]
dtype: object
user_32169800
Related