Pandas slice string with index from another column

Question

I want to slice this string using indexes from another column. I'm getting NaN instead of slices of the string.

import pandas as pd
from pandas import DataFrame, Series

sales = {'name': ['MSFTCA', 'GTX', 'MSFTUSA', ],
         'n_chars': [2, 2, 3],
         'Jan': [150, 200, 50],
         'Feb': [200, 210, 90],
         'Mar': [140, 215, 95]}
df = pd.DataFrame.from_dict(sales)
df

def extract_location(name, n_chars):
    return( name.str[-n_chars:])

df.assign(location=(lambda x: extract_location(x['name'], x['n_chars']))).to_dict()

Gives:

{'Feb': {0: 200, 1: 210, 2: 90},
 'Jan': {0: 150, 1: 200, 2: 50},
 'Mar': {0: 140, 1: 215, 2: 95},
 'location': {0: nan, 1: nan, 2: nan},
 'n_chars': {0: 2, 1: 2, 2: 3},
 'name': {0: 'MSFTCA', 1: 'GTX', 2: 'MSFTUSA'}}

Answer 1

You need apply with axis=1 for processing by rows:

def extract_location(name, n_chars):
    return( name[-n_chars:])


df=df.assign(location=df.apply(lambda x: extract_location(x['name'], x['n_chars']), axis=1))
print (df) 
   Feb  Jan  Mar  n_chars     name location
0  200  150  140        2   MSFTCA       CA
1  210  200  215        2      GTX       TX
2   90   50   95        3  MSFTUSA      USA

---

df = df.assign(location=df.apply(lambda x: x['name'][-x['n_chars']:], axis=1))
print (df) 
   Feb  Jan  Mar  n_chars     name location
0  200  150  140        2   MSFTCA       CA
1  210  200  215        2      GTX       TX
2   90   50   95        3  MSFTUSA      USA

Answer 2

Using a comprehension

df.assign(location=[name[-n:] for n, name in zip(df.n_chars, df.name)])

   Feb  Jan  Mar  n_chars     name location
0  200  150  140        2   MSFTCA       CA
1  210  200  215        2      GTX       TX
2   90   50   95        3  MSFTUSA      USA

You can speed it up a bit with

df.assign(location=[name[-n:] for n, name in zip(df.n_chars.values, df.name.values)])