replace value if string ends with certain value pandas dataframe

Viewed 767

As titled.

I have written my code but it does not work. I wish I can get a more pythonic way of writing the code (in one single line perhaps).

clean_df :

columnA
 123F
 FVGD
 w999Z
 678Q
 6y6yA

My code :

postfix = ["A", "D", "Z", "P"]

for value in postfix:
    if cleaned_data['columnA'].str.endswith(value) is True:
        cleaned_data['columnA'] = value
    else:
        cleaned_data['columnA'] = "blah"

The postfix are constant. Expected outcome :

columnA
 blah
  D
  Z
 blah
  A
4 Answers

In one line with a list comprehension:

postfix = ["A", "D", "Z", "P"]
cleaned_data['columnA'] = [value[-1] if value[-1] in postfix else "blah" for value in cleaned_data['columnA']]

The output is :

columnA
 blah
  D
  Z
 blah
  A

You can try this with pd.Series.str.extract here.

pat = "|".join(postfix)
pat = f"({pat}$"
df['columnA'] = df['columnA'].str.extract(pat, expand=False).fillna('blah')
df
  columnA
0    blah
1       D
2       Z
3    blah
4       A

A simple one-liner would be

df['columnA'] = np.where(df.columnA.str[-1].isin(postfix), df.columnA.str[-1], 'blah')

np.where takes in a condition, value if the condition is True and value if the condition is False.

OR

In pure pandas, without using numpy, it would be

df['columnA'] = df.columnA.str[-1].where(df.columnA.str[-1].isin(postfix), 'blah')

You can use numpy.where:

Note: str.endswith accepts a tuple also.

In [3933]: import numpy as np

In [3934]: df.columnA = np.where(df.columnA.str.endswith(tuple(postfix)), df.columnA.str[-1], 'blah')

In [3935]: df
Out[3935]: 
  columnA
0    blah
1       D
2       Z
3    blah
4       A
Related