Pandas - New column based on the value of another column N rows back, when N is stored in a column

Viewed 484

I have a pandas dataframe with example data:

idx    price    lookback
0      5    
1      7        1
2      4        2
3      3        1
4      7        3
5      6        1

Lookback can be positive or negative but I want to take the absolute value of it for how many rows back to take the value from.

I am trying to create a new column that contains the value of price from lookback + 1 rows ago, for example:

idx    price    lookback  lb_price
0      5        NaN       NaN
1      7        1         NaN
2      4        2         NaN
3      3        1         7
4      7        3         5
5      6        1         3

I started with what felt like the most obvious way, this did not work:

df['sbc'] = df['price'].shift(dataframe['lb'].abs() + 1)

I then tried using a lambda, this did not work but I probably did it wrong:

sbc = lambda c, x: pd.Series(zip(*[c.shift(x+1)]))
df['sbc'] = sbc(df['price'], df['lb'].abs())

I also tried a loop (which was extremely slow, but worked) but I am sure there is a better way:

lookback = np.nan
for i in range(len(df)):
    if df.loc[i, 'lookback']:
        if not np.isnan(df.loc[i, 'lookback']):
            lookback = abs(int(df.loc[i, 'lookback']))
    if not np.isnan(lookback) and (lookback + 1) < i:
        df.loc[i, 'lb_price'] = df.loc[i - (lookback + 1), 'price']

I have seen examples using lambda, df.apply, and perhaps Series.map but they are not clear to me as I am quite a novice with Python and Pandas.

I am looking for the fastest way I can do this, if there is a way without using a loop.

Also, for what its worth, I plan to use this computed column to create yet another column, which I can do as follows:

df['streak-roc'] = 100 * (df['price'] - df['lb_price']) / df['lb_price']

But if I can combine all of it into one really efficient way of doing it, that would be ideal.

Solution!

Several provided solutions worked great (thank you!) but all needed some small tweaks to deal with my potential for negative numbers and that it was a lookback + 1 not - 1 and so I felt it was prudent to post my modifications here.

All of them were significantly faster than my original loop which took 5m 26s to process my dataset.

I marked the one I observed to be the fastest as accepted as I improving the speed of my loop was the main objective.

Edited Solutions

From Manas Sambare - 41 seconds

df['lb_price'] = df.apply(
    lambda x: df['price'][x.name - (abs(int(x['lookback'])) + 1)]
    if not np.isnan(x['lookback']) and x.name >= (abs(int(x['lookback'])) + 1)
    else np.nan,
    axis=1)

From mannh - 43 seconds

def get_lb_price(row, df):
    if not np.isnan(row['lookback']):
        lb_idx = row.name - (abs(int(row['lookback'])) + 1)
        if lb_idx >= 0:
            return df.loc[lb_idx, 'price']
        else: 
            return np.nan

df['lb_price'] = dataframe.apply(get_lb_price, axis=1 ,args=(df,))

From Bill - 18 seconds

lookup_idxs = df.index.values - (abs(df['lookback'].values) + 1)
valid_lookups = lookup_idxs >= 0
df['lb_price'] = np.nan
df.loc[valid_lookups, 'lb_price'] = df['price'].to_numpy()[lookup_idxs[valid_lookups].astype(int)]
4 Answers

By getting the row's index inside of the df.apply() call using row.name, you can generate the 'lb_price' data relative to which row you are currently on.

%time
df.apply(
    lambda x: df['price'][x.name - int(x['lookback'] + 1)]
    if not np.isnan(x['lookback']) and x.name >= x['lookback'] + 1 
    else np.nan,
    axis=1)

# > CPU times: user 2 µs, sys: 0 ns, total: 2 µs
# > Wall time: 4.05 µs

FYI: There is an error in your example as idx[5]'s lb_price should be 3 and not 7.

Here is an example which uses a regular function

def get_lb_price(row, df):
    lb_idx = row.name - abs(row['lookback']) - 1
    if lb_idx >= 0:
        return df.loc[lb_idx, 'price']
    else: 
        return np.nan

df['lb_price'] = df.apply(get_lb_price, axis=1 ,args=(df,))

This solution loops of the values ot the column lockback and calculates the index of the wanted value in the column price which I store as a list.

The rule it, that the lockback value has to be a number and that the wanted index is not smaller than 0.

new = np.zeros(df.shape[0])
price = df.price.values
for i, lookback in enumerate(df.lookback.values):
    # lookback has to be a number and the index is not allowed to be less than 0
    # 0<i-lookback is equivalent to 0<=i-(lookback+1)
    if lookback!=np.nan and 0<i-lookback: 
        new[i] = price[int(i-(lookback+1))]
    else:
        new[i] = np.nan
df['lb_price'] = new

Here's a vectorized version (i.e. no for loops) using numpy array indexing.

lookup_idxs = df.index.values - df['lookback'].values - 1
valid_lookups = lookup_idxs >= 0
df['lb_price'] = np.nan
df.loc[valid_lookups, 'lb_price'] = df.price.to_numpy()[lookup_idxs[valid_lookups].astype(int)]
print(df)

Output:

     price  lookback  lb_price
idx                           
0        5       NaN       NaN
1        7       1.0       NaN
2        4       2.0       NaN
3        3       1.0       7.0
4        7       3.0       5.0
5        6       1.0       3.0
Related