Python: return dictionary of keys for columns, values for lists of row indexes where the cell is equ-CodePudding

Say I have a dataframe like so:

pd.DataFrame({'Col1': [0, 1],
              'Col2': [1, 1],
              'Col3': [0, 0]}, index=[1, 2])

Index	Col1	Col2	Col3
1	0	1	0
2	1	1	0

I want to return a dictionary, where each column is a key, and each value is a list of row indexes where the value of row col is equal to 1. The output of the following dataframe I want is:

    {'Col1': [2],
     'Col2': [1, 2]
     'Col3': []}

I could probably figure it out myself using some for loops but I suppose there must be an easier, possibly built-in method.

CodePudding user response：

You can compare all values by 1 for boolean DataFrame and then create dictionary in dict comprehension with filter index values:

d = {k:df.index[v].tolist() for k, v in df.eq(1).items()}
print (d)
{'Col1': [2], 'Col2': [1, 2], 'Col3': []}

Another idea with DataFrame.agg:

d = df.eq(1).agg(lambda x: x.index[x].tolist()).to_dict()
print (d)
{'Col1': [2], 'Col2': [1, 2], 'Col3': []}