如何用Pandas对指定列范围应用apply生成符合条件的列头列表？

原始数据 #

我们有如下DataFrame：

| North | South | East | West |
|-------|-------|------|------|
| 8     | 1     | 8    | 6    |
| 4     | 4     | 8    | 4    |
| 1     | 1     | 1    | 2    |
| 7     | 3     | 7    | 8    |

需求 #

新增Headers列，仅针对动态选择的列子集（比如示例中的South、East列），记录每行中值为1的对应列名列表，预期输出：

Headers 
|---------------|
| [South]       |
|               |
| [South, East] |
|               |

错误代码问题分析 #

你尝试的两种写法都存在问题，具体如下：

第一种错误写法 #

df['Headers'] = df.iloc[N1:N2].apply(lambda x: df.columns[x==1].tolist(),axis=1)

1. iloc用法错误：df.iloc[N1:N2]是提取行子集而非列子集，要取列子集需写成df.iloc[:, N1:N2]（冒号表示所有行）。
2. 列名不匹配：即使修正了iloc的写法，lambda里用df.columns调用的是全量列名，而非你选中的列子集的列名，会导致匹配到错误的列名。

第二种错误写法 #

df['Headers'] = df.apply(lambda x: df.iloc[N1:N2].columns[x==1].tolist(),axis=1)

1. 布尔数组长度不匹配：lambda中的x是原DataFrame的整行数据，x==1生成的布尔数组长度等于原DataFrame的总列数，但df.iloc[N1:N2].columns是列子集的列名列表，长度仅为N2-N1，两者长度不一致，会引发索引错误。
2. 同样存在iloc取行的问题：这里df.iloc[N1:N2]还是取行，不是列子集，应该改为df.iloc[:, N1:N2]。

正确实现方法 #

方法一：先提取列子集再处理 #

先单独取出目标列子集，再对其应用逻辑，这样列名和行数据完全对应：

# 先获取列子集，N1、N2为列的索引范围（比如取第1到第3列，对应South、East）
subset = df.iloc[:, N1:N2]
# 对列子集每行筛选值为1的列名
df['Headers'] = subset.apply(lambda x: subset.columns[x == 1].tolist(), axis=1)

方法二：在apply中直接遍历目标列 #

如果不想单独提取子集，也可以在lambda里直接遍历目标列，判断对应值是否为1：

# 获取目标列的列名列表
target_cols = df.columns[N1:N2]
# 遍历每行，收集值为1的列名
df['Headers'] = df.apply(lambda x: [col for col in target_cols if x[col] == 1], axis=1)

两种方法都能得到你预期的结果。

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