Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Get name of colums where true in pandas dataframe

Tags:

python

pandas

Suppose I have the following pandas dataframe:

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

And I want, for each id, get the name of the columns where True, the final dict would be:

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

I think maybe it is possible with a group by replacing True with the name of the column and them aggregating it with list but I couldn't come up with a solution.

EDIT: If it is all false for an id, then just return an empty list.

like image 867
Bruno Mello Avatar asked Jan 24 '23 14:01

Bruno Mello


1 Answers

Try with melt then 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']}
like image 98
BENY Avatar answered May 22 '23 15:05

BENY