Vinícius Aguiar Vinícius Aguiar - 11 months ago 61
Python Question

Get list with element's columns from Pandas DataFrame

I need to have a list containing all specific element's columns for every index. For example, this DataFrame:

>>> df
1 2 3 4 5
2016-01-27 A B B I I
2016-03-07 A C D U U
2016-04-12 H A V V V
2016-05-02 B L Y S N
2016-05-23 L N N A S

Inputting "A" I'd like to have this list as output:


Is there a built-in method for this?

Edit: In the original table all items in a row are unique, when editing original table to make it less "dense" to post here and I made this mistake, sorry.

Answer Source

You may want to melt your data frame to long format and then calculate the corresponding list of columns for each input(value), After obtaining the Series as follows, it would be easy for you to query the result for any intended input:

import pandas as pd

#A    [1, 1, 2, 4]
#B       [1, 2, 3]
#C             [2]
#D             [3]
#H             [1]
#I          [4, 5]
#L          [1, 2]
#N       [2, 3, 5]
#S          [4, 5]
#U          [4, 5]
#V       [3, 4, 5]
#Y             [3]
#Name: variable, dtype: object

To get the list of columns for input A:

result = pd.melt(df).groupby('value').variable.apply(list)

# ['1', '1', '2', '4']