Pandas如何根据值非零条件将对应列名赋值给行的新增列

需求说明 #

针对DataFrame的每一行，筛选取值不为0的列，将这些列的列名用英文逗号拼接后，赋值给新增的Col_name字段。

实现代码 #

基础实现（易读性优先） #

import pandas as pd

# 测试数据构造
df = pd.DataFrame({
    'date' : ['2021-08-01', '2021-08-01', '2021-08-01', '2021-08-02', '2021-08-02', '2021-08-02'],
    'person': ['type_A', 'type_C', 'type_C', 'type_A', 'type_B', 'type_D'],
    'London' : [1, 4, 0, 5, 7, 9],
    'New York' : [0.0, 0.0, 3, 0.0, 0.0, 0.0],
    'Boston' : [3, 6, 9, 1, 0.0, 7],
    'Hong Kong' : [0.0, 0.0, 0.0, 0.0, 3, 0.0]
})

# 指定需要校验非0值的列，排除无需参与判断的date、person字段
check_columns = ['London', 'New York', 'Boston', 'Hong Kong']

# 逐行筛选非0值对应的列名，拼接后赋值
df['Col_name'] = df[check_columns].apply(lambda row: ', '.join(row[row != 0].index), axis=1)

print(df)

高效实现（大数量场景优先） #

如果处理的DataFrame数据量较大，逐行apply的性能较低，可以使用矩阵运算的方式提速：

# 布尔矩阵点积实现，避免逐行遍历，性能提升显著
df['Col_name'] = df[check_columns].ne(0).dot(check_columns + ',').str.rstrip(',')

输出结果 #

运行上述代码后得到的结果符合逻辑要求：

date  person  London  New York  Boston  Hong Kong          Col_name
0  2021-08-01  type_A       1       0.0     3.0        0.0     London, Boston
1  2021-08-01  type_C       4       0.0     6.0        0.0     London, Boston
2  2021-08-01  type_C       0       3.0     9.0        0.0   New York, Boston
3  2021-08-02  type_A       5       0.0     1.0        0.0     London, Boston
4  2021-08-02  type_B       7       0.0     0.0        3.0  London, Hong Kong
5  2021-08-02  type_D       9       0.0     7.0        0.0     London, Boston

注：你给出的预期输出中第5行Col_name多了Hong Kong，属于和原始数据冲突的笔误，上述输出为符合逻辑的正确结果。

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