首页 > 解决方案 > 从 pandas Dataframe 列创建列表的高效/Pythonic 方式

问题描述

我有一个dataframe如下。

df = pd.DataFrame({
    'code' : [1,   2,  3,  4,  5,  6,  7,  8,  9,  10],
    'Tag' :  ['A','B','C','D','B','C','D','A','D','C']
})

+------+-----+
| code | Tag |
+------+-----+
|   1  |  A  |
+------+-----+
|   2  |  B  |
+------+-----+
|   3  |  C  |
+------+-----+
|   4  |  D  |
+------+-----+
|   5  |  B  |
+------+-----+
|   6  |  C  |
+------+-----+
|   7  |  D  |
+------+-----+
|   8  |  A  |
+------+-----+
|   9  |  D  |
+------+-----+
|  10  |  C  |
+------+-----+

code我的目标是根据列中的常见项目创建列表Tag,如下所示。

codes_A = [1,8]
codes_B = [2,5]
codes_C = [3,6,10]
codes_D = [4,7,9]

我现在的做法是

codes_A = df[df['Tag'] == 'A']['code'].to_list()
codes_B = df[df['Tag'] == 'B']['code'].to_list()
codes_C = df[df['Tag'] == 'C']['code'].to_list()
codes_D = df[df['Tag'] == 'D']['code'].to_list()

这段代码完成了这项工作。但是,正如您所看到的,这非常麻烦且效率低下。我多次重复相同的代码,并且在我想创建新列表时也重复。

在or中是否有更有效的pythonic方法来做到这一点?pandasnumpy

标签: pythonpandasnumpy

解决方案


创建列表字典,因为不推荐使用变量名

d = df.groupby('Tag')['code'].agg(list).to_dict()
print (d)
{'A': [1, 8], 'B': [2, 5], 'C': [3, 6, 10], 'D': [4, 7, 9]}

然后通过dict中的键进行列表查找,但没有分配给变量名:

print (d['A'])
[1, 8]

所以实际上它意味着在你的代码中如果使用codes_A然后它会更改为d['A'],对于所有变量都类似。


But if really need it:

for k, v in d.items():
    globals()[f'code_{k}'] = v
    
print (code_A)
[1, 8]

推荐阅读