CODEWITHSUNDEEP
pythonpandas
AmirX
AmirX

Reputation: 2717

Remove suffix from certain column names in Pandas

This is my df:

df = pd.DataFrame({'a_x':[1, 2, 3], 'b_x':[20, 30, 40], 'c_x':[10, 40, 50]})

How can I remove the suffix of first and third column names?

I don't want this approach: df = df.rename(columns={'a_x':'a', 'c_x':'c'}) that hardcoded them one by one.

EDIT 1: I have the list of the columns that I want to remove the suffix from. In this case I have ['a', 'c']

My desired outcome looks like this:

    a    b_x   c
0    1   20   10
1    2   30   40
2    3   40   50

Upvotes: 1

Views: 642

Answers (3)

sammywemmy
sammywemmy

Reputation: 28729

This works for python3.8 with the walrus operator; since you already have a list of the columns:

df.rename(columns = lambda x: val if (val:=x.split("_")[0]) in ["a", "c"] else x)

   a  b_x   c
0  1   20  10
1  2   30  40
2  3   40  50

Upvotes: 1

Umar.H
Umar.H

Reputation: 23099

I think a simple if/else and a list comprehension will do here.

import pandas as pd 

trg_cols = ['a_x','c_x']
new_cols = [col.split('_')[0] if col in trg_cols else col  for col in df.columns]

df.columns = new_cols 
print(df)
   a  b_x   c
0  1   20  10
1  2   30  40
2  3   40  50

or if you want to select by your output target columns.

trg_cols = ['a','c']
new_cols = [col if not col.split('_')[0] in trg_cols else col.split('_')[0] 
            for col in df.columns]


df.columns = new_cols
print(df)
   a  b_x   c
0  1   20  10
1  2   30  40
2  3   40  50

of usingfillna and pd.Series

trg_cols = ['a','c']
s = pd.Series(df.columns)
df.columns = s.str.extract('('+'|'.join(trg_cols)+')',expand=False).fillna(s)


print(df)
   a  b_x   c
0  1   20  10
1  2   30  40
2  3   40  50

Upvotes: 6

Sayandip Dutta
Sayandip Dutta

Reputation: 15872

You can define your cols, and rename:

>>> target_cols = ['c']
>>> suffix = '_x'
>>> df.rename(columns={col + suffix: col for col in target_cols})

   a_x  b_x   c
0    1   20  10
1    2   30  40
2    3   40  50

# or
# df.rename(columns=dict(zip((col + suffix for col in target_cols), target_cols)))

Upvotes: 5

Related Questions