join column names in a new pandas columns conditional on value

user_9906395 0
Viewed 362

I have the following dataset:

data = {'Environment': ['0', '0', '0'],
        'Health': ['1', '0', '1'],
            'Labor': ['1', '1', '1'],
             }

df = pd.DataFrame(data, columns=['Environment', 'Health', 'Labor'])

I want to create a new column df['Keyword'] whose value is a join of the column names with value > 0.

Expected Outcome:

data = {'Environment': ['0', '0', '0'],
            'Health': ['1', '0', '1'],
                'Labor': ['1', '1', '1'],
                     'Keyword': ['Health, Labor', 'Labor', 'Health, Labor']}
    
df_test = pd.DataFrame(data, columns=['Environment', 'Health', 'Labor', 'Keyword']) 
df_test
df = pd.DataFrame(data, columns=['Environment', 'Health', 'Labor'])

How do I go about it?

pandas join python
3 Answers

Other version with .apply():

df['Keyword'] = df.apply(lambda x: ', '.join(b for a, b in zip(x, x.index) if a=='1'),axis=1)
print(df)

Prints:

  Environment Health Labor        Keyword
0           0      1     1  Health, Labor
1           0      0     1          Labor
2           0      1     1  Health, Labor
user_100359850

Another method with mask and stack then groupby to get your aggregation of items.

stack by default drops na values.

df['keyword'] = df.mask(
               df.lt(1)).stack().reset_index(1)\
                        .groupby(level=0)["level_1"].agg(list)

print(df)

   Environment  Health  Labor          keyword
0            0       1      1  [Health, Labor]
1            0       0      1          [Labor]
2            0       1      1  [Health, Labor]
user_93751020

First problem in sample data values are strings, so if want compare for greater use:

df = df.astype(float).astype(int)

Or:

 df = df.replace({'0':0, '1':1})

And then use DataFrame.dot for matrix multiplication with columns names and separators, last remove it from right side:

df['Keyword'] = df.gt(0).dot(df.columns + ', ').str.rstrip(', ')
print (df)
   Environment  Health  Labor        Keyword
0            0       1      1  Health, Labor
1            0       0      1          Labor
2            0       1      1  Health, Labor

Or compare strings - e.g. not equal '0' or equal '1':

df['Keyword'] = df.ne('0').dot(df.columns + ', ').str.rstrip(', ')

df['Keyword'] = df.eq('1').dot(df.columns + ', ').str.rstrip(', ')
user_29010020
Related