Как добавить список имен столбцов в DataFrame в pandas при условиях

Question

У меня есть DataFrame из pandas, который показывает некоторые слова в качестве имен столбцов. Значение в ячейках - это количество раз, когда слово используется в письме (строке):

+--------+-------+-------+-------+-------+
| index  | word1 | word2 | word3 | word4 |
+--------+-------+-------+-------+-------+
|      0 |     1 |     2 |     1 |     0 |
|      1 |     2 |     3 |     5 |     1 |
|      2 |     0 |     0 |     1 |     0 |
+--------+-------+-------+-------+-------+

Теперь мне нужен список со словами (именами столбцов) в конце каждой строки, но только если слово используется. что-то вроде этого:

+--------+-------+-------+-------+-------+---------------------------+
| index  | word1 | word2 | word3 | word4 |           text            |
+--------+-------+-------+-------+-------+---------------------------+
|      0 |     1 |     2 |     1 |     0 | [word1,word2,word3]       |
|      1 |     2 |     3 |     5 |     1 | [word1,word2,word3,word4] |
|      2 |     0 |     0 |     3 |     0 | [word3]                   |
+--------+-------+-------+-------+-------+---------------------------+

Я знаю, что могу получить список с list(data.columns), но я не получаю, как поставить условия и добавить новый столбец со списком в нем.

Answer 1

In [136]: df = pd.DataFrame(np.random.randint(0, 3, (3, 5)), columns=list('abcde'))

In [137]: df
Out[137]:
   a  b  c  d  e
0  1  0  1  0  1
1  0  2  0  0  2
2  0  1  1  0  0

In [140]: df['text'] = df.apply(lambda x: df.columns[x.astype(bool)].to_list(), axis=1)

In [141]: df
Out[141]:
   a  b  c  d  e       text
0  1  0  1  0  1  [a, c, e]
1  0  2  0  0  2     [b, e]
2  0  1  1  0  0     [b, c]