Строка извлекает имена столбцов из DataFrame в Series

Question

Я бы хотел извлечь имена столбцов в списке в Серию, отфильтрованную по значениям в каждой строке

In [1]: import pandas as pd   

In [2]: df =pd.DataFrame({'colA':[1,0,1], 'colB':[0,0,1], 'colC':[1,0,0]})    

In [3]: print(df)

   colA  colB  colC
0     1     0     1
1     0     0     0
2     1     1     0

Получившийся ряд должен выглядеть следующим образом:

0    [colA, colC]
1              []
2    [colA, colB]
dtype: object

Вот измученное решение, которое я придумал:

In [4]: df2 = df.T

In [5]: l = [df2[df2[i]>0].index.values.tolist() for i in range(3)]

In [6]: print(pd.Series(l))

0    [colA, colC]
1              []
2    [colA, colB]
dtype: object

Есть ли менее замученный способ сделать это?

Answer 1

Вы можете использовать np.where, предполагая, что ваш фрейм данных состоит из 0 и 1, и создать серию из результата:

x = np.where(df,df.columns,'')
pd.Series([' '.join(i).split() for i in x])
0    [colA, colC]
1              []
2    [colA, colB]