CODEWITHSUNDEEP
pythonpandasnumpy
Franc Weser
Franc Weser

Reputation: 869

Access last X rows in Pandas apply

I have a Pandas dataframe with one column of numbers, similar to this:

id - val
0  - 100
1  - 200
2  - 100
3  - 400
4  - 300
5  - 100
etc

What I would like to do is to add a second column which is a list/numpy array of the values from the previous 3 rows:

id - val - val_list
0  - 100 - [] # Or [NaN, NaN, NaN]
1  - 200 - [100] # Or [NaN, NaN, 100]
2  - 100 - [100, 200] # Or [NaN, 100, 200]
3  - 400 - [100, 200, 100]
4  - 300 - [200, 100, 400]
5  - 100 - [100, 400, 300]
etc

Any idea how to solve this efficiently, preferably without looping?

Upvotes: 4

Views: 832

Answers (2)

Bhanu Tez
Bhanu Tez

Reputation: 306

Hi use the following simple code.

df = pd.DataFrame([100,200,100,400,300,100],columns =['Val'])

temp = pd.concat([df.shift(3),df.shift(2),df.shift(1)],axis=1)
df['val_list'] = temp.apply(lambda x:x.tolist(),axis=1)
# 
df = pd.DataFrame([100,200,100,400,300,100],columns =['Val'])
N=3
temp = pd.DataFrame()
for i in range(N,0,-1):
    temp = pd.concat([temp,df.shift(i)],axis=1)
df['val_list'] = temp.apply(lambda x:x.tolist(),axis=1)

Upvotes: 4

jezrael
jezrael

Reputation: 862611

First I think working with lists in pandas is not good idea, if possible better is working with 2d numpy array here.

Use strides if performance is important:

N = 3
x = np.concatenate([[np.nan] * (N), df['val'].values])

def rolling_window(a, window):
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
arr = rolling_window(x, N)

df['val_list'] = arr[:-1].tolist()
print (df)
   id  val               val_list
0   0  100        [nan, nan, nan]
1   1  200      [nan, nan, 100.0]
2   2  100    [nan, 100.0, 200.0]
3   3  400  [100.0, 200.0, 100.0]
4   4  300  [200.0, 100.0, 400.0]
5   5  100  [100.0, 400.0, 300.0]

Upvotes: 5

Related Questions