Pandas apply to a range of columns

Question

Given the following dataframe, I would like to add a fifth column that contains a list of column headers when a certain condition is met on a row, but only for a range of dynamically selected columns (ie subset of the dataframe)


| North | South | East | West |
|-------|-------|------|------|
| 8     | 1     | 8    | 6    |
| 4     | 4     | 8    | 4    |
| 1     | 1     | 1    | 2    |
| 7     | 3     | 7    | 8    |

For instance, given that the inner two columns ('South', 'East') are selected and that column headers are to be returned when the row contains the value of one (1), the expected output would look like this:

Headers 
|---------------|
| [South]       |
|               |
| [South, East] |
|               |

The following one liner manages to return column headers for the entire dataframe.

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

I tried adding the dynamic column range condition by using iloc but to no avail. What am I missing?

For reference, these are my two failed attempts (N1 and N2 being column range variables here)

df['Headers'] = df.iloc[N1:N2].apply(lambda x: df.columns[x==1].tolist(),axis=1)

df['Headers'] = df.apply(lambda x: df.iloc[N1:N2].columns[x==1].tolist(),axis=1)

Nev1111 · Answer

This works:

df=pd.DataFrame({'North':[8,4,1,7],'South':[1,4,1,3],'East':[8,8,1,7],\
                 'West':[6,4,2,8]})


df1=df.melt(ignore_index=False)

condition1=df1['variable']=='South'
condition2=df1['variable']=='East'
condition3=df1['value']==1
    
df1=df1.loc[(condition1|condition2)&condition3]
df1=df1.groupby(df1.index)['variable'].apply(list)
df=df.join(df1)

Pandas apply to a range of columns

1 Answers