首页 > 解决方案 > 如何从一组列创建映射字典

问题描述

我有一个如下所示的 Pandas DataFrame:

import numpy as np
import pandas as pd

data = np.random.rand(18).reshape(-1, 6)
data = pd.DataFrame(data, columns = ['var1_x10', 'var2_x10', 'var3_x10', 'var1_x20', 'var2_x20', 'var3_x20'])

    var1_x10    var2_x10    var3_x10    var1_x20    var2_x20    var3_x20
0   0.171464    0.441099    0.936246    0.532478    0.128823    0.211489
1   0.917217    0.544899    0.589996    0.362159    0.774122    0.439542
2   0.094015    0.582171    0.573968    0.200833    0.257705    0.057575

如您所见,列实际上是每个原始列var1var2、的 2 个转换var3。现在我想以字典的形式创建一个映射,其中原始列名作为键,转换后的列名列表作为值:

my_dict = {'var1': ['var1_x10', 'var1_x20'], 
           'var2': ['var2_x10', 'var2_x20'], 
           'var3': ['var3_x10', 'var3_x20']}

我怎样才能做到这一点?

标签: pythonpandas

解决方案


使用Series.groupbywith 将列转换为系列split并转换为list

d = data.columns.to_series().groupby(lambda x: x.split('_')[0]).apply(list).to_dict()
print (d)
{'var1': ['var1_x10', 'var1_x20'], 
 'var2': ['var2_x10', 'var2_x20'], 
 'var3': ['var3_x10', 'var3_x20']}

另一种解决方案:

from collections import defaultdict
d = defaultdict (list)

for x in data.columns:
    d[x.split('_')[0]].append(x)

print (dict(d))
{'var1': ['var1_x10', 'var1_x20'], 
 'var2': ['var2_x10', 'var2_x20'], 
 'var3': ['var3_x10', 'var3_x20']}

推荐阅读