Enumerate columns with same prefix
The idea is to group columns with the same prefix, and establish a cumcount for them.
Since we need to handle column without a prefix separately, we will need to do this in two steps using GroupBy.cumcount
and np.where
:
cols = df.columns.str.split('_').str[0].to_series()
df.columns = np.where(
cols.groupby(level=0).transform('count') > 1,
cols.groupby(level=0).cumcount().add(1).astype(str).radd(df.columns),
cols
)
df
A B Data_mean1 Data_std2 Data_corr3 Text_one1 Text_two2 Text_three3
0 a e 1 5 9 foo bar bar
1 b f 2 6 10 bar foo bar
2 c g 3 7 11 foobar barfoo barbar
3 d h 4 8 12 barfoo foobar foofoo
A simpler solution would be to set columns you don't want to add a suffix to as the index. Then you can simply do
df.set_index(['A', 'B'], inplace=True)
df.columns = (
df.columns.str.split('_')
.str[0]
.to_series()
.groupby(level=0)
.cumcount()
.add(1)
.astype(str)
.radd(df.columns))
df
Data_mean1 Data_std2 Data_corr3 Text_one1 Text_two2 Text_three3
A B
a e 1 5 9 foo bar bar
b f 2 6 10 bar foo bar
c g 3 7 11 foobar barfoo barbar
d h 4 8 12 barfoo foobar foofoo
You could also use a defaultdict to create a counter for each prefix.
from collections import defaultdict
prefix_starting_location = 2
columns = df.columns[prefix_starting_location:]
prefixes = set(col.split('_')[0] for col in columns)
new_cols = []
dd = defaultdict(int)
for col in columns:
prefix = col.split('_')[0]
dd[prefix] += 1
new_cols.append(col + str(dd[prefix]))
df.columns = df.columns[:prefix_starting_location].tolist() + new_cols
>>> df
A B Data_mean1 Data_std2 Data_corr3 Text_one1 Text_two2 Text_three3
0 a e 1 5 9 foo bar bar
1 b f 2 6 10 bar foo bar
2 c g 3 7 11 foobar barfoo barbar
3 d h 4 8 12 barfoo foobar foofoo
If the prefixes are known:
prefixes = ['Data', 'Text']
new_cols = []
dd = defaultdict(int)
for col in df.columns:
prefix = col.split('_')[0]
if prefix in prefixes:
dd[prefix] += 1
new_cols.append(col + str(dd[prefix]))
else:
new_cols.append(col)
If your split character _
is not in any of your data fields:
new_cols = []
dd = defaultdict(int)
for col in df.columns:
if '_' in col:
prefix = col.split('_')[0]
dd[prefix] += 1
new_cols.append(col + str(dd[prefix]))
else:
new_cols.append(col)
df.columns = new_cols
you can use rename
such as:
l_word = ['Data','Text']
df = df.rename(columns={ col:col+str(i+1)
for word in l_word
for i, col in enumerate(df.filter(like=word))})