pandanic way of inserting df[col].str.extract() results back into original Pandas df immediately after the extraction column

user_6226980 0
Viewed 95

Please refer to the below runnable demo code. It is trying to sr.str.extract() the "a" column into multiple columns, and insert those columns into the origianl df immediately after the "a" column.
Please combine steps [1] and [2] in the below code in a better way.

import re
import pandas as pd

df = pd.DataFrame({
     "a" : {1: 'a', 2: 'aa', 3: 'aaa'} ,
     "b" : {1: 'b', 2: 'bb', 3: 'bbb'} ,
     "c" : {1: 'b', 2: 'bb', 3: 'bbb'} ,
})

"""
df==
+----+-----+-----+-----+
|    | a   | b   | c   |
|----+-----+-----+-----|
|  1 | a   | b   | b   |
|  2 | aa  | bb  | bb  |
|  3 | aaa | bbb | bbb |
+----+-----+-----+-----+
"""
# step [1] sr.str.extract
rex = re.compile(r'(?P<firstletter>\w) (?P<secondletter>\w+)', re.X)
cols = df['a'].str.extract(rex)

# step [2] insert extracted columns back into the original df
df['firstletter'], df['secondletter'] = 0, 0
df['firstletter'] = cols['firstletter']
df['secondletter'] = cols['secondletter']
df = df['a firstletter secondletter b c'.split()]

"""
# Or, a more concise step [2], but too hard to glance thru and remember, also prone to mistake:
for col in cols.columns[::-1]:
    df.insert(df.columns.get_loc('a')+1, col, cols[col])
"""

# result:
"""
df==
+----+-----+---------------+----------------+-----+-----+
|    | a   | firstletter   | secondletter   | b   | c   |
|----+-----+---------------+----------------+-----+-----|
|  1 | a   | nan           | nan            | b   | b   |
|  2 | aa  | a             | a              | bb  | bb  |
|  3 | aaa | a             | aa             | bbb | bbb |
+----+-----+---------------+----------------+-----+-----+
"""
pandas python regex
2 Answers

There are 2 possible solutions:

Solution 1:

This solution would most resemble your ideal solution while recommend you make some slight (but constructive) change to the naming of the extracted columns. E.g. instead of firstletter and secondletter, for the extraction of column a, we name it with prefix a_ as a_firstletter and a_secondletter. Then we can use one statement as follows:

import re

df = df.assign(**df['a'].str.extract(r'(?P<a_firstletter>\w) (?P<a_secondletter>\w+)', re.X)).sort_index(axis=1)

Result:

print(df)


     a a_firstletter a_secondletter    b    c
1    a           NaN            NaN    b    b
2   aa             a              a   bb   bb
3  aaa             a             aa  bbb  bbb

Solution 2:

This solution you can keep on using firstletter and secondletter as the extracted column names.

You can split the columns into 2 portions by .iloc : left portion df_left from first column till column a and right portion df_right from the columns after column a till the end. Then concat the left portion df_left, the newly extracted columns cols, and right portion df_right together along the columns by pd.concat(), as follows:

df_left = df.iloc[:, 0: df.columns.get_loc('a')+1]
df_right = df.iloc[:, df.columns.get_loc('a')+1:]

df = pd.concat([df_left, cols, df_right], axis=1)

Result:

print(df)


     a firstletter secondletter    b    c
1    a         NaN          NaN    b    b
2   aa           a            a   bb   bb
3  aaa           a           aa  bbb  bbb
user_150706970

I believe you need.

import re

df = pd.DataFrame({
     "a" : {1: 'a', 2: 'aa', 3: 'aaa'} ,
     "b" : {1: 'b', 2: 'bb', 3: 'bbb'} ,
     "c" : {1: 'b', 2: 'bb', 3: 'bbb'} ,
})

df[['firstletter','secondletter']] = df['a'].str.extract(r"(\w)(\w+)", expand=True)
print(df)

Output:

     a    b    c firstletter secondletter
1    a    b    b         NaN          NaN
2   aa   bb   bb           a            a
3  aaa  bbb  bbb           a           aa
user_5323120
Related