在python中查找相应的列

Question

我有一个数据集simalar到这个：

p = {'A': [0,1,0,1], 'B': [1,1,1,1], 'C': [0,0,1,1], 'D': [1,1,1,0]}
df5 = pd.DataFrame(data=p)
df5

现在我想创建一个列表，每行对应一个值，我目前正在这样做：

cols = df5.dot(df5.columns).map(set).values.tolist()
cols

但是，如果我的专栏名称不是“A”，而是“AA”，那么这就不再有效了，是否有解决方法

Answer 1

您可以为列名添加分隔符，然后通过^{}删除last并使用^{}：

p = {'AA': [0,1,0,1], 'B': [1,1,1,1], 'C': [0,0,1,1], 'D': [1,1,1,0]}
df5 = pd.DataFrame(data=p)

cols = df5.dot(df5.columns + ',').str.rstrip(',').str.split(',').map(set).values.tolist()
print (cols)
[{'D', 'B'}, {'B', 'D', 'AA'}, {'C', 'D', 'B'}, {'B', 'AA', 'C'}]

另一种解决方案是使用numpy索引：

c = df5.columns.to_numpy()
cols = [set(c[x]) for x in df5.to_numpy().astype(bool)]
print (cols)
[{'D', 'B'}, {'B', 'D', 'AA'}, {'C', 'D', 'B'}, {'B', 'AA', 'C'}]