python - Pandas DataFrame 聚合列，其他列的名称作为值

Question

我正在尝试在我的 DataFrame 中创建一个新列，这是一个聚合列名列表。这是一个示例 DataFrame:

In [1]: df = pd.DataFrame({'A':[1,2,3],
               'B':[4,5,6],
               'C':[7,8,9],
               'D':[1,3,5],
               'E':[5,3,6],
               'F':[7,4,3]})
In [2]: df
Out[2]:
   A  B  C  D  E  F
0  1  4  7  1  5  7
1  2  5  8  3  3  4
2  3  6  9  5  6  3

我想创建一个新列，其中包含满足特定条件的列名列表。假设我对值 > 3 的列感兴趣——我想要一个如下所示的输出:

In [3]: df
Out[3]:
   A  B  C  D  E  F  Flag
0  1  4  7  1  5  7  ['B', 'C', 'E', 'F']
1  2  5  8  3  3  4  ['B', 'C', 'F']
2  3  6  9  5  6  3  ['B', 'C', 'D', 'E']

目前，我正在使用apply:

df['Flag'] = df.apply(lambda row: [list(df)[i] for i, j in enumerate(row) if j > 3], axis = 1)

这样可以完成工作，但感觉很笨拙，我想知道是否有更优雅的解决方案。

谢谢!

Answer 1

使用df.dot()这里:

df['Flag']=(df>3).dot(df.columns).apply(list)
print(df)

---

   A  B  C  D  E  F          Flag
0  1  4  7  1  5  7  [B, C, E, F]
1  2  5  8  3  3  4     [B, C, F]
2  3  6  9  5  6  3  [B, C, D, E]