在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

将1值替换为列名

df5.replace(1, pd.Series(df5.columns, df5.columns), inplace=True)

将0值替换为NaN，然后使用stack删除它们并转换为list

cols = df5.replace(0, np.nan).stack().groupby(level=0).apply(list).tolist()
cols

这将返回列表列表，而不是集合列表：

[['B', 'D'], ['A', 'B', 'D'], ['B', 'C', 'D'], ['A', 'B', 'C']]

堆叠是为了去除零。如果您可以保留它们（或者在不同的步骤中删除），那么df5.values.tolist()就可以了

Answer 2

您可以为列名添加分隔符，然后通过^{}删除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'}]