Перебор строк датафрейма в Pandas (Python)

Question

Необходимо написать функцию, которая будет принимать каждую строку датафрейма, и возвращать названия столбцов (со значением 1) в виде списка. Допустим, есть датафрейм:

import pandas as pd
inp = [{'name':'ab', 'A':1, 'B':1}, {'name':'cd','A':0, 'B':1}, {'name':'ef','A':1, 'B':1}]
df = pd.DataFrame(inp)
print(df)

name   A  B
ab     1  1
cd     0  1
ef     1  1

Нужно написать функцию, чтобы она принимала каждую строку и возвращала названия столбцов со значением 1 в виде списка. У меня уже есть отдельная переменная columns = ['A', 'B'], содержащая все названия столбцов кроме первого в виде списка.

В идеале для первой строки должен быть выход list = ['A', 'B'], так как оба столбца первой строчки равны единице.

Что у меня получилось:

def func(columns):
    for i in columns:
        list=[]
        for j in df[i]:
            if j == 1:
                list.append(i)
            return list

Но получается, что она каждый раз возвращает названия столбцов отдельными списками. Понимаю, что порядок изначально не тот, но не могу придумать, как поменять местами список columns и df[i]. В итоге забрел в тупик.

Answer 1

попробуйте так:

In [119]: cols = ["A", "B"]

In [120]: a = np.array(cols)

In [121]: res = df[cols].eq(1).apply(lambda row: a[row].tolist(), axis=1)

результат в виде Pandas.Series:

In [122]: res
Out[122]:
0    [A, B]
1       [B]
2    [A, B]
dtype: object

в виде DataFrame:

In [123]: df.join(res.to_frame(name="res"))
Out[123]:
  name  A  B     res
0   ab  1  1  [A, B]
1   cd  0  1     [B]
2   ef  1  1  [A, B]

БЛОГ НА HUSL

Перебор строк датафрейма в Pandas (Python)

Ответы (1 шт):