python pandas - как создать для каждой строки список имен столбцов с условием? - PullRequest
Новая
       71

python pandas - как создать для каждой строки список имен столбцов с условием?

1 голос
/ 01 мая 2020

Мне нужно применить функцию ко всем строкам кадра данных. Я использовал эту функцию, которая возвращает список имен столбцов, если значение равно 1: 

def find_column(x):  
    a=[]  
    for column in df.columns:  
        if (df.loc[x,column] == 1):  
            a = a + [column]
    return a

, это работает, если я просто вставляю индекс, например : 

print(find_column(1))

но: 

df['new_col'] = df.apply(find_column,axis=1)

не работает ни одной идеи? Спасибо!

1 Ответ

0 голосов
/ 01 мая 2020

Вы можете выполнять итерацию по каждой строке, поэтому x равно Series с index, такими же, как имена столбцов, поэтому возможна фильтрация сопоставленных данных индекса и преобразование в список: 

df = pd.DataFrame({
        'A':list('abcdef'),
         'B':[4,1,4,5,5,1],
         'C':[7,1,9,4,2,3],
         'D':[1,1,5,7,1,1],
         'E':[5,1,6,9,1,4],
         'F':list('aaabbb')
})

def find_column(x):
    return x.index[x == 1].tolist()

df['new'] = df.apply(find_column,axis=1)
print (df)
   A  B  C  D  E  F           new
0  a  4  7  1  5  a           [D]
1  b  1  1  1  1  a  [B, C, D, E]
2  c  4  9  5  6  a            []
3  d  5  4  7  9  b            []
4  e  5  2  1  1  b        [D, E]
5  f  1  3  1  4  b        [B, D]

Другая идея используйте DataFrame.dot с маской на DataFrame.eq для равных, затем удалите последний разделитель и используйте Series.str.split: 

df['new'] = df.eq(1).dot(df.columns + ',').str.rstrip(',').str.split(',')
print (df)

   A  B  C  D  E  F           new
0  a  4  7  1  5  a           [D]
1  b  1  1  1  1  a  [B, C, D, E]
2  c  4  9  5  6  a            []
3  d  5  4  7  9  b            []
4  e  5  2  1  1  b        [D, E]
5  f  1  3  1  4  b        [B, D]
...