Python实现DataFrame中值为1的列序号数组生成并添加为新列

嗨，这个需求很容易实现，用Pandas就能轻松搞定！我给你一步步拆解：

首先，我们先把你的示例数据构造出来：

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'ID': [0, 1, 2, 3],
    'Q1': [1, 0, 1, 0],
    'Q2': [0, 0, 1, 1],
    'Q3': [1, 1, 1, 1]
})

接下来，我们要生成Array列，核心思路是针对Q1、Q2、Q3这三列，逐行找出值为1的列的索引（也就是你示例里的[0,2]这种格式）：

# 定义要处理的目标列
target_cols = ['Q1', 'Q2', 'Q3']

# 逐行生成值为1的列索引数组
df['Array'] = df[target_cols].apply(
    lambda row: [idx for idx, value in enumerate(row) if value == 1],
    axis=1  # axis=1表示按行处理
)

# 调整列的顺序，让Array列放在ID后面，和你想要的格式一致
df = df[['ID', 'Array'] + target_cols]

运行完这段代码后，打印df就能得到你想要的结果：

ID       Array  Q1  Q2  Q3
0   0     [0, 2]   1   0   1
1   1        [2]   0   0   1
2   2  [0, 1, 2]   1   1   1
3   3     [1, 2]   0   1   1

补充说明： #

• 用target_cols变量指定目标列，后续如果列名有变化，直接修改这个列表就行，非常灵活。
• apply方法配合lambda函数的写法很直观，适合小到中等数据量的场景；如果你的数据量特别大，可以考虑用numpy的向量化操作来提升效率，比如：

import numpy as np

# 用numpy实现更高效的版本
arr = df[target_cols].values == 1
df['Array'] = [np.where(row)[0].tolist() for row in arr]

这个方法的速度会比apply快很多，适合百万级别的数据行。

备注：内容来源于stack exchange，提问作者Haris Jawed