为Dataframe中重复的指定列名添加后缀

问题场景 #

现有一个Pandas DataFrame，列名依次为week、SITE、LAL、SITE、LAL，需要将所有名为SITE的列添加数字后缀，最终列名变为week、SITE_1、LAL、SITE_2、LAL。

示例DataFrame：

week    SITE        LAL SITE        LAL
0   1   BARTON CHAPEL   1.1 PENASCAL I  1.0
1   2   BARTON CHAPEL   1.1 PENASCAL I  1.0
2   3   BARTON CHAPEL   1.1 PENASCAL I  1.0
3   4   BARTON CHAPEL   1.1 PENASCAL I  1.0
4   5   BARTON CHAPEL   1.1 PENASCAL I  1.0
5   6   BARTON CHAPEL   1.4 PENASCAL I  1.0

解决方案 #

可以借助itertools.count生成递增的数字后缀，遍历列名完成重命名：

步骤1：构造示例DataFrame（若已有可跳过）

import pandas as pd

df = pd.DataFrame(
    [
        [1, 'BARTON CHAPEL', 1.1, 'PENASCAL I', 1.0],
        [2, 'BARTON CHAPEL', 1.1, 'PENASCAL I', 1.0],
        [3, 'BARTON CHAPEL', 1.1, 'PENASCAL I', 1.0],
        [4, 'BARTON CHAPEL', 1.1, 'PENASCAL I', 1.0],
        [5, 'BARTON CHAPEL', 1.1, 'PENASCAL I', 1.0],
        [6, 'BARTON CHAPEL', 1.4, 'PENASCAL I', 1.0],
    ],
    columns=['week', 'SITE', 'LAL', 'SITE', 'LAL']
)

步骤2：重命名列

from itertools import count

# 初始化计数器，从1开始
site_counter = count(1)
# 遍历列名，为SITE列添加后缀
df.columns = [f"{col}_{next(site_counter)}" if col == 'SITE' else col for col in df.columns]

验证结果 #

执行后查看列名：

print(df.columns)

输出：

Index(['week', 'SITE_1', 'LAL', 'SITE_2', 'LAL'], dtype='object')

完全符合需求。

内容的提问来源于stack exchange，提问作者user2100039