如何高效将Pandas DataFrame中的虚拟变量转为分类列？

嘿，这个循环效率低的问题我太熟了！别担心，用pandas的向量化操作分分钟解决，比逐行循环快好几个数量级，完全不用再写慢到离谱的循环啦。

方法一：用idxmax直接生成（最简洁高效） #

你的数据是标准的虚拟变量格式——每行只有一个1，对应所属的站点ID。idxmax方法可以沿着行方向找到最大值（也就是1）所在的列名，直接把它赋值给新的Sites列，这是pandas内部优化的向量化操作，全程在C层面执行，速度碾压循环。

示例代码：

import pandas as pd

# 假设你的原始DataFrame名为df
# 先指定要处理的虚拟变量列（确保只选这三列）
site_cols = ['ID01', 'ID18', 'ID31']
# 生成Sites分类列
df['Sites'] = df[site_cols].idxmax(axis=1)

如果你的数据里可能存在全0的行（没有任何站点标记），可以用fillna给这些行设置默认值：

df['Sites'] = df[site_cols].idxmax(axis=1).fillna('Unknown')

方法二：用melt重塑数据（适合复杂场景） #

如果后续需要对站点数据做更多重塑操作，也可以用melt把宽表转成窄表，再筛选出值为1的行：

# 假设你的DataFrame还有其他需要保留的列，比如时间列，把它们放在id_vars里
melted_df = df.melt(
    id_vars=['其他需要保留的列名'],  # 替换成你实际需要保留的列
    value_vars=site_cols,
    var_name='Sites',
    value_name='is_selected'
)
# 只保留标记为1的行，得到最终的DataFrame
final_df = melted_df[melted_df['is_selected'] == 1].drop('is_selected', axis=1)

为什么循环这么慢？ #

Python的for循环逐行处理DataFrame是出了名的低效——因为pandas的底层是numpy数组，向量化操作是在C语言层面执行的，而Python循环是在解释器层面逐行处理，两者的效率差能达到几十甚至上百倍。所以只要能用上pandas的内置向量化方法，就千万别写循环！

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