Issue

This Content is from Stack Overflow. Question asked by gmorissette

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)

Solution

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)

This Question was asked in StackOverflow by gmorissette and Answered by Nev1111 It is licensed under the terms of CC BY-SA 2.5. - CC BY-SA 3.0. - CC BY-SA 4.0.