基于列表从数据框中提取值作为字典

时间:2018-02-20 06:46:11

标签: python pandas

我在每列中都有一个具有唯一值的数据框:

df1 = pd.DataFrame([["Phys","Shane","NY"],["Chem","Mark","LA"],
                    ["Maths","Jack","Mum"],["Bio","Sam","CT"]],
                    columns = ["cls1","cls2","cls3"])
print(df1)

    cls1    cls2    cls3
0   Phys    Shane   NY
1   Chem    Mark    LA
2   Maths   Jack    Mum
3   Bio     Sam     CT

列表l1:

l1=["Maths","Bio","Shane","Mark"]
print(l1)

['Maths', 'Bio', 'Shane', 'Mark']

现在我想从数据框中检索包含列表和元素列表中元素的列。

预期输出

{'cls1' : ['Maths','Bio'], 'cls2': ['Shane','Mark']}

我的代码

cls = []
for cols in df1.columns:
    mask = df1[cols].isin(l1)
    if mask.any():
        cls.append(cols)
print(cls)

以上代码的输出

['cls1', 'cls2']

我正在努力从数据框和列表中获取常用元素以将其转换为字典。

欢迎任何建议。

感谢。

1 个答案:

答案 0 :(得分:2)

使用DataFrame.isin作为掩码,通过索引替换非匹配值并重新整形为stack

df = df1[df1.isin(l1)].stack()
print (df)
0  cls2    Shane
1  cls2     Mark
2  cls1    Maths
3  cls1      Bio
dtype: object

dict comprehension的最后一个创建列表:

d = {k:v.tolist() for k,v in df.groupby(level=1)}
print(d)
{'cls2': ['Shane', 'Mark'], 'cls1': ['Maths', 'Bio']}

另一种解决方案:

d = {}
for cols in df1.columns:
    mask = df1[cols].isin(l1)
    if mask.any():
        d[cols] = df1.loc[mask, cols].tolist()

print(d)
{'cls2': ['Shane', 'Mark'], 'cls1': ['Maths', 'Bio']}