如何获取pandas get_dummies生成的哑变量列名？

当使用pandas.get_dummies()处理DataFrame的分类列后，若想快速获取生成的哑变量列名，可采用以下几种简洁高效的实现方式，替代原有的繁琐循环写法：

方法一：提前生成列名（无需遍历生成后的DataFrame） #

这种方式直接基于原分类列的唯一值生成预期的哑变量列名，不需要处理生成后的DataFrame列，效率更高：

import pandas as pd

# 构建示例DataFrame
df = pd.DataFrame({
    "numeric_1": [1, 2, 3, 4],
    "cat_1": ['high', 'low', 'high', 'low'],
    "cat_2": ['blue', 'blue', 'green', 'green']
})

categorical_colnames = ['cat_1', 'cat_2']

# 提前生成哑变量列名
dummy_cols = []
for col in categorical_colnames:
    # 获取当前分类列的唯一值
    unique_vals = df[col].unique()
    # 拼接列名并添加到列表
    dummy_cols.extend([f"{col}_{val}" for val in unique_vals])

# 生成哑变量
df = pd.get_dummies(df, columns=categorical_colnames)

print(dummy_cols)
# 输出: ['cat_1_high', 'cat_1_low', 'cat_2_blue', 'cat_2_green']

方法二：从生成后的DataFrame中提取列名 #

如果已经完成哑变量生成，可通过列表推导式简化提取逻辑：

简化版循环写法 #

# 生成哑变量后直接提取
dummy_cols = [col for col in df.columns if any(col.startswith(f"{cat}_") for cat in categorical_colnames)]

print(dummy_cols)
# 输出: ['cat_1_high', 'cat_1_low', 'cat_2_blue', 'cat_2_green']

更高效的前缀匹配写法 #

利用str.startswith()支持元组参数的特性，进一步简化代码：

# 构造前缀元组
prefixes = tuple(f"{cat}_" for cat in categorical_colnames)
# 匹配并提取列名
dummy_cols = [col for col in df.columns if col.startswith(prefixes)]

print(dummy_cols)
# 输出: ['cat_1_high', 'cat_1_low', 'cat_2_blue', 'cat_2_green']

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