获取pandas数据帧中为true的列的名称

Question

假设我有以下pandas数据帧：

df = pd.DataFrame({'id': [1,2,3,4,5], 'a': [True, True, False, True, False], 'b': [False, True, False, False, True], 'c': [False, False, True, True, True]})
id      a      b      c
1   True  False  False
2   True   True  False
3  False  False   True
4   True  False   True
5  False   True   True

我想，对于每个id，得到True所在列的名称，最后的dict是：

{1: ['a'], 2: ['a', 'b'], 3: ['c'], 4: ['a', 'c'], 5: ['b', 'c']}

我认为也许可以通过将True替换为列的名称并将其聚合为list来创建组，但我无法想出解决方案。

编辑：如果一个id都是false，那么只返回一个空列表。

Answer 1

尝试使用melt，然后使用groupby

out_d = df.melt('id').query('value').groupby('id')['variable'].agg(list).to_dict()
Out[127]: {1: ['a'], 2: ['a', 'b'], 3: ['c'], 4: ['a', 'c'], 5: ['b', 'c']}

Answer 2

另一种方法是好的(？？？(旧apply，然后是to_dict

df.set_index('id').apply(lambda x: list(x[x].index), axis=1).to_dict()

输出：

{1: ['a'], 2: ['a', 'b'], 3: ['c'], 4: ['a', 'c'], 5: ['b', 'c']}

Answer 3

您可以尝试将id设置为索引、堆叠、仅筛选True值，然后使用defaultdict获取字典：

box = df.set_index("id").stack().loc[lambda x: x].index
from collections import defaultdict
d = defaultdict(list)
for k, v in box:
d[k].append(v)
d
defaultdict(list,
{1: ['a'], 2: ['a', 'b'], 3: ['c'], 4: ['a', 'c'], 5: ['b', 'c']})

我估计它可能比@BEN_YO的绝妙解决方案快一点。和往常一样，您的里程数可能会有所不同。

如果id全部为False：

df = pd.DataFrame(
{
"id": [1, 2, 3, 4, 5],
"a": [True, True, False, True, False],
"b": [False, True, False, False, False],
"c": [False, False, True, True, False],
}
)
df
id     a       b    c
0   1   True    False   False
1   2   True    True    False
2   3   False   False   True
3   4   True    False   True
4   5   False   False   False

运行上面的代码后，您可以执行额外的步骤：

{k: d[k] for k in df.id}
{1: ['a'], 2: ['a', 'b'], 3: ['c'], 4: ['a', 'c'], 5: []}