如何为DataFrame动态生成列，将每行值为Y的value类列对应数值转为列表？

给定如下格式的DataFrame：

import pandas as pd

df = pd.DataFrame({'column1': ['Y', 'Y', 'Y'],
                   'value_5': ['N', 'Y', 'Y'],
                   'value_6': ['N', 'Y', 'N'],
                   'value_10': ['Y', 'N', 'N'],
                   'value_20': ['N', 'N', 'Y']},
                  index=['key1','key2','key4'])

输出的DataFrame如下：

column1 value_5 value_6 value_10 value_20
key1       Y       N       N        Y        N
key2       Y       Y       Y        N        N
key4       Y       Y       N        N        Y   

需求：为该DataFrame新增最后一列，每行的值为对应行中值为Y的value_*类列名的后缀数字组成的列表。注意：每次运行时，DataFrame的列数量和值可能存在差异。

可以通过以下步骤实现：

• 筛选目标列：先提取所有以value_开头的列，这一步能自动适配不同数量的目标列。
• 逐行处理：对每行数据，找出值为Y的列，提取列名中的后缀数字，组成列表。

具体代码如下：

# 筛选出所有value_开头的列
value_cols = [col for col in df.columns if col.startswith('value_')]

# 定义函数处理每行数据
def get_y_numbers(row):
    # 筛选当前行值为Y的列，提取后缀数字并转为整数
    return [int(col.split('_')[1]) for col in value_cols if row[col] == 'Y']

# 新增结果列
df['result'] = df.apply(get_y_numbers, axis=1)

运行后得到的DataFrame如下：

column1 value_5 value_6 value_10 value_20    result
key1       Y       N       N        Y        N      [10]
key2       Y       Y       Y        N        N   [5, 6]
key4       Y       Y       N        N        Y      [5, 20]

代码说明 #

• value_cols = [col for col in df.columns if col.startswith('value_')]：动态获取所有目标列，不管列数量多少都能适配。
• get_y_numbers函数：遍历目标列，判断当前行该列是否为Y，如果是则拆分列名取后缀数字，转为整数后加入列表。
• df.apply(..., axis=1)：按行应用函数，生成结果列。

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