Check substring in dataframe column by referencing string position

user_11609503 0
Viewed 364

I have a dataframe with a column called 'id', where each id is 9 characters long, and I'm trying to add a column 'Rating' where I rate each row as A, AA, or AAA based on whether the 6th, 7th and 8th characters are 'A00', 'AA0' or '000'. I've got the following code so far:

id = df['id']
conditions = [(id.str.get(5) == 'A00'), (id.str.get(5) == 'AA0'), (id.str.get(5) == '000')]
values = ['A', 'AA', 'AAA']
df['Rating'] = np.select(conditions, values)
df['Rating'] = df['Rating'].astype('category')

But I know the conditions line is wrong because the column.str.get(n) only checks the (n-1)th character of the string and I need the substring of length 3 rather than just a single character. Does anyone know which command I can use?

Thanks in advance!

pandas substring indexing python string
2 Answers

Using replace() with value mappings.:

df.id.str[-4:-1].replace({'A00': 'A', 'AA0': 'AA', '000': 'AAA'})
user_34424090

Just extract it using a regular expression.

df.id.str.extract(r'.{5}([A]+)').fillna('AAA')

Example

df = pd.DataFrame({'id': ['12345A003', '12345AA03', '123450003']}) 
df.id.str.extract(r'.{5}([A]+)').fillna('AAA')

Output

     0
0    A
1   AA
2  AAA
user_138133650
Related