CODEWITHSUNDEEP
pythonpandas
Roy Han
Roy Han

Reputation: 141

Adding conditional prefixes to column names

So I have a dataframe with some weird suffixes, like _a or _b that map to certain codes, I was wondering how would you add a prefix depending on the the suffix and remove the suffix to something easier to understand.

i.e.:

    red_a    blue_a    green_b
....
....

to 

    A red    A blue   B green
....
....

I tried 

for col in df.columns:
    if col.endswith('_a'):
        batch_match[col].replace('_a', '')
        batch_match[col].add_prefix('A ')
    else:
        batch_match[col].add_prefix('B ')

But it returns a df of NaN.

Upvotes: 0

Views: 889

Answers (2)

pault
pault

Reputation: 43504

You can use pandas.DataFrame.rename with a customer mapper

df = pd.DataFrame(
    {"red_a": ['a', 'b', 'c'], "blue_a": [1, 2, 3], 'green_b': ['x', 'y', 'z']}
)

def renamer(col):
    if any(col.endswith(suffix) for suffix in ['_a', '_b']):
        prefix = col[-1]                        # use last char as prefix
        return prefix.upper() + " " + col[:-2]  # add prefix and strip last 2 chars
    else:
        return col

df = df.rename(mapper=renamer, axis='columns')
print(df)
#   A blue B green A red
#0       1       x     a
#1       2       y     b
#2       3       z     c

Upvotes: 1

BENY
BENY

Reputation: 323236

What I will do 

df.columns=df.columns.str.split('_').map(lambda x : '{} {}'.format(x[1].upper(),x[0]))
df
Out[512]: 
  A red  A blue B green
0     a       1       x
1     b       2       y
2     c       3       z

Upvotes: 0

Related Questions