pandas printing a column if columns value is 1 applied to all columns

user_1587457 0
Viewed 275

I have rows and columns with columns representing actual entities. The Column values apart from the first column are either 1 or 0. The first column is a key. The objective is to return the column name (2nd to last column) if the column value is 1. This is the function that i have written and it works. Was wondering if there is a better way to express this in Pandas, or even a better way to represent this form of data to make it more pandas friendly.

def return_keys(df,productname):
    df2 = df[df['Product']==productname]
    print(df2)
    columns = list(df2)
    cust=[]
    for col in columns[1:]:
        if (df2[col].to_list()[0]==1):
            cust.append(col)
    return cust
pandas python
2 Answers

If your key column does not contain 0/1 , you can try using apply row-wise. Below is an example dataset:

import pandas as pd
import numpy as np

np.random.seed(111)
df = pd.DataFrame({'Product':np.random.choice(['A','B','C'],10),
'Col1':np.random.binomial(1,0.5,10),
'Col2':np.random.binomial(1,0.5,10),
'Col3':np.random.binomial(1,0.5,10)})

df

  Product  Col1  Col2  Col3
0       A     0     1     1
1       A     1     0     0
2       A     1     1     1
3       A     1     0     0
4       C     1     1     1
5       B     0     1     1
6       C     1     0     0
7       C     0     1     0
8       C     1     1     1
9       A     0     1     0

We apply a boolean and apply (axis=1) onto this boolean data.frame, call out the columns.

(df == 1).apply(lambda x:df.columns[x].tolist(),axis=1)

0          [Col2, Col3]
1                [Col1]
2    [Col1, Col2, Col3]
3                [Col1]
4    [Col1, Col2, Col3]
5          [Col2, Col3]
6                [Col1]
7                [Col2]
8    [Col1, Col2, Col3]
9                [Col2]
user_122584590

Try the following:

df = pd.DataFrame({'a':[1,2,3], 'b':[0,1,2], 'c':[1,2,4], 'd':[0,2,4]})
cols_with_first_element_1 = df.columns[df.iloc[0]==1].to_list()
print(cols_with_first_element_1)

results in ['a', 'c'].

user_66584220
Related