Apply customized function for extracting numbers from string to multiple columns in Python

user_8410477 0
Viewed 245

Given a dataset as follows:

   id name   value1   value2 value3
0   1   gz        1        6     st
1   2   sz  7-tower  aka 278      3
2   3   wh      NaN       67   s6.1
3   4   sh       x3       45     34

I'd like to write a customized function to extract numbers from values columns.

Here is the pseudocode I have written:

def extract_nums(row):
    return row.str.extract('(\d*\.?\d+)', expand=True)

df[['value1', 'value2', 'value3']] = df[['value1', 'value2', 'value3']].apply(extract_nums)

It raise an error:

ValueError: If using all scalar values, you must pass an index

Code for manipulating columns one by one works, but not concise:

df['value1'] = df['value1'].str.extract('(\d*\.?\d+)', expand=True)
df['value2'] = df['value2'].str.extract('(\d*\.?\d+)', expand=True)
df['value3'] = df['value3'].str.extract('(\d*\.?\d+)', expand=True)

How could write the code correctly? Thanks.

pandas dataframe python-3.x regex
2 Answers

You can filter the value like columns then stack the columns and use str.extract to extract the numbers followed by unstack to reshape:

c = df.filter(like='value').columns
df[c] = df[c].stack().str.extract('(\d*\.?\d+)', expand=False).unstack()

Alternatively you can try str.extract with apply:

c = df.filter(like='value').columns
df[c] = df[c].apply(lambda s: s.str.extract('(\d*\.?\d+)', expand=False))

Result:

   id name value1 value2 value3
0   1   gz      1      6    NaN
1   2   sz      7    278      3
2   3   wh    NaN     67    6.1
3   4   sh      3     45     34
user_128331660

The following code also works out:

cols = ['value1', 'value2', 'value3']
for col in cols:
    df[col] = df[col].str.extract('(\d*\.?\d+)', expand=True)

Alternative solution with function:

def extract_nums(row):
    return row.str.extract('(\d*\.?\d+)', expand=False)

df[cols] = df[cols].apply(extract_nums)

Out:

   id name value1 value2 value3
0   1   gz      1      6    NaN
1   2   sz      7    278      3
2   3   wh    NaN     67    6.1
3   4   sh      3     45     34
user_84104770
Related