如何为Pandas DataFrame生成符合条件的列名列表列，规避循环与赋值错误？

现有如下结构的Pandas DataFrame，需要生成choices列：收集每行中值为1的列名后缀（如choice_a取a）组成列表，若没有值为1的列则设为NA。原循环代码运行报错，需解决错误并尽可能避免使用循环。

示例数据：

data = [[0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 0, 1], [0, 0, 0]]
df = pd.DataFrame(data, columns=["choice_a", "choice_b", "choice_c"])

预期结果：

choice_a  choice_b  choice_c  choices
index                  
0        0       0           1         ['c']
1        0       1           0         ['b']
2        1       0           0         ['a']
3        1       0           1         ['a','c']
4        0       0           0         NA

原报错代码：

df['choices']=0
for i in np.arange(df.shape[0]):
    choice_list = []
    for j in np.arange(len(df.columns)):
        if df.iloc[i,j]==1:
            choice_list.append(df.columns[j].split('_')[1])
    df.iloc[i,4]=choice_list

报错信息：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/65/3mqr9fpn37jf2xt2pxbcgp_w0000gn/T/ipykernel_1513/2279334138.py in <module>
      5         if main_dataset.iloc[i,j]==1:
      6             choice_list.append(main_dataset.columns[j].split('_')[1])
----> 7     main_dataset.iloc[i,5]=choice_list

~/opt/anaconda3/lib/python3.9/site-packages/pandas/core/indexing.py in __setitem__(self, key, value)
    714 
    715         iloc = self if self.name == "iloc" else self.obj.iloc
--> 716         iloc._setitem_with_indexer(indexer, value, self.name)
    717 
    718     def _validate_key(self, key, axis: int):

~/opt/anaconda3/lib/python3.9/site-packages/pandas/core/indexing.py in _setitem_with_indexer(self, indexer, value, name)
   1689         if take_split_path:
   1690             # We have to operate column-wise
-> 1691             self._setitem_with_indexer_split_path(indexer, value, name)
   1692         else:
   1693             self._setitem_single_block(indexer, value, name)

~/opt/anaconda3/lib/python3.9/site-packages/pandas/core/indexing.py in _setitem_with_indexer_split_path(self, indexer, value, name)
   1744                     return self._setitem_with_indexer((pi, info_axis[0]), value[0])
   1745 
-> 1746                 raise ValueError(
   1747                     "Must have equal len keys and value "
   1748                     "when setting with an iterable"

ValueError: Must have equal len keys and value when setting with an iterable

报错原因：直接通过iloc给单个单元格赋值列表时，Pandas会默认将列表视为要分配给多个元素的序列，而非单个单元格的值，因此触发长度不匹配的错误。

以下是两种无需嵌套循环的高效实现方式：

方法一：使用apply+列表推导式（代码简洁，适合小数据集） #

import pandas as pd
import numpy as np

data = [[0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 0, 1], [0, 0, 0]]
df = pd.DataFrame(data, columns=["choice_a", "choice_b", "choice_c"])

# 提前提取所有列名的后缀
col_suffixes = [col.split('_')[1] for col in df.columns]

# 逐行筛选值为1的列对应的后缀，生成列表
df['choices'] = df.apply(
    lambda row: [col_suffixes[i] for i, val in enumerate(row) if val == 1],
    axis=1
)

# 将空列表替换为NA
df['choices'] = df['choices'].map(lambda x: x if x else np.nan)

方法二：使用melt+groupby（完全矢量化，适合大数据集） #

import pandas as pd
import numpy as np

data = [[0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 0, 1], [0, 0, 0]]
df = pd.DataFrame(data, columns=["choice_a", "choice_b", "choice_c"])

# 重置索引以便后续分组
df_reset = df.reset_index()
# 宽表转长表，拆分列名与对应值
melted = df_reset.melt(id_vars='index', var_name='choice_col', value_name='value')
# 提取列名后缀
melted['suffix'] = melted['choice_col'].str.split('_').str[1]
# 按索引分组，收集值为1的后缀成列表
choices = melted[melted['value'] == 1].groupby('index')['suffix'].agg(list)
# 合并回原DataFrame，无匹配的行自动填充NA
df = df.join(choices.rename('choices'))

两种方法都能得到预期结果，其中方法二完全基于Pandas的矢量化操作，避免了逐行处理，在数据量较大时性能更优。

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