Get column name based on condition in pandas

Question

I have a dataframe as below:

I want to get the name of the column if column of a particular row if it contains 1 in the that column.

Answer 1

Getting column name are dividing in 2 sections.

If you want in a new column name then condition should be unique because it will only give 1 col name for each row.

data = {'foo':[0,0,3,0], 'bar':[0, 5, 0, 0], 'baz':[0,0,2,0], 'spam':[0,1,0,1]}
df = pd.DataFrame(data)
df=df.replace(0,np.nan)
df
    foo bar baz spam
0   NaN NaN NaN NaN
1   NaN 5.0 NaN 1.0
2   3.0 NaN 2.0 NaN
3   NaN NaN NaN 1.0

If you were looking for min or maximum

max= df.idxmax(1)
min = df.idxmin(1)
out=  df.assign(max=max , min=min)
out

    foo bar baz spam   max    min
0   NaN NaN NaN NaN    NaN    NaN
1   NaN 5.0 NaN 1.0    bar    spam
2   3.0 NaN 2.0 NaN    foo    baz
3   NaN NaN NaN 1.0    spam   spam

2nd case, If your condition is satisfied in multiple columns for example you are looking for columns that contain 1 and you are looking for list because its not possible to adjust in same dataframe.

str_con= df.astype(str).apply(lambda x:x.str.contains('1.0',case=False, na=False)).any()
df.column[str_con]
#output
Index(['spam'], dtype='object') #only spam contains 1

Or you are looking for numerical condition columns contains value more than 1

num_con = df.apply(lambda x:x>1.0).any()
df.columns[num_con]
#output
Index(['foo', 'bar', 'baz'], dtype='object') #these col has higher value than 1

Happy learning

Answer 2

Firstly

Your question is very ambiguous and I recommend reading this link in @sammywemmy's comment. If I understand your problem correctly... we'll talk about this mask first:

df.columns[      
    (df == 1)        # mask 
    .any(axis=0)     # mask
]

What's happening? Lets work our way outward starting from within df.columns[**HERE**] :

1. (df == 1) makes a boolean mask of the df with True/False(1/0)
2. .any() as per the docs:

"Returns False unless there is at least one element within a series or along a Dataframe axis that is True or equivalent".

This gives us a handy Series to mask the column names with.

We will use this example to automate for your solution below

---

Next:

Automate to get an output of (<row index> ,[<col name>, <col name>,..]) where there is 1 in the row values. Although this will be slower on large datasets, it should do the trick:

import pandas as pd

data = {'foo':[0,0,0,0], 'bar':[0, 1, 0, 0], 'baz':[0,0,0,0], 'spam':[0,1,0,1]}
df = pd.DataFrame(data, index=['a','b','c','d'])

print(df)

   foo  bar  baz  spam
a    0    0    0     0
b    0    1    0     1
c    0    0    0     0
d    0    0    0     1

# group our df by index and creates a dict with lists of df's as values
df_dict = dict(
    list(
        df.groupby(df.index)
    )
)

Next step is a for loop that iterates the contents of each df in df_dict, checks them with the mask we created earlier, and prints the intended results:

for k, v in df_dict.items():               # k: name of index, v: is a df
    check = v.columns[(v == 1).any()]
    if len(check) > 0:
        print((k, check.to_list()))

('b', ['bar', 'spam'])
('d', ['spam'])

Side note:

You see how I generated sample data that can be easily reproduced? In the future, please try to ask questions with posted sample data that can be reproduced. This way it helps you understand your problem better and it is easier for us to answer it for you.

Answer 3

Use DataFrame.dot:

df1 = df.dot(df.columns)

If there is multiple 1 per row:

df2 = df.dot(df.columns + ';').str.rstrip(';')