python - pandas dataframe print column if value = 1

Question

Have this data in panda dataframe:

animal      tall    swim        jump  
frog          0        1            1
toad          1        0            0
tadpole       0        0            1

Would like this output:

frog swim
frog jump
toad tall
tadpole jump

Any builtin function for this?

This should get you close: `df.set_index('animal').dot(df.columns[1:] + ',').str.strip(',')` — cs95, Aug 03 '18 at 07:06
@coldspeed it is close, frog swim,jump is one one line though — tooptoop4, Aug 03 '18 at 10:42
@jezrael can u reopen ? idxmax will only return 1 result per line, we want all that have 1 in a line — tooptoop4, Aug 03 '18 at 10:45

jezrael · Accepted Answer · 2018-08-03T11:19:57.273

Use set_index with stack, filter 1 and create DataFrame by contructor:

df1 = df.set_index('animal').stack()
s = df1[df1==1]

df2 = pd.DataFrame({'a':s.index.get_level_values(0),
                    'b':s.index.get_level_values(1)})
print (df2)
         a     b
0     frog  swim
1     frog  jump
2     toad  tall
3  tadpole  jump

Another solution is use MultiIndex.to_frame:

df2 = s.index.to_frame(index=False)
df2.columns = ['a','b']
print (df2)
         a     b
0     frog  swim
1     frog  jump
2     toad  tall
3  tadpole  jump

Different solution with melt and filtering by query:

df2 = (df.melt('animal', var_name='val')
         .query('value == 1')
         .sort_values('animal')
         .reset_index(drop=True)
         .drop('value', axis=1))

python - pandas dataframe print column if value = 1

1 Answers1