根据 Pandas DF 中的每行条件获取列标题列表

Question

I was wondering if it were possible to get a list of column headers based on a condition. For example, if the condition I have is to get a list of the column headers that had a "MATCH" value in each cell, it would output either a list of lists or a list of strings containing the header name, as such:

["a, c", "b, d", "a, b, c, d", "a, d"]
or 
[["a", "c"], ["b", "d"], ["a", "b", "c", "d"], ["a", "d"]]

Thank you for any help!

Answer 1

You could try with np.where :

import pandas as pd
import numpy as np
df=pd.DataFrame({'a': ['match','mismatch','match'],'b': ['match','match','mismatch'],'c': ['mismatch','mismatch','match']})

print(df)

arr= np.where(df.eq('match'), df.columns, '').sum(axis=1)

print(arr)

Output:

df
          a         b         c
0     match     match  mismatch
1  mismatch     match  mismatch
2     match  mismatch     match

arr
['ab' 'b' 'ac']

And then, to get the desired lists you could try:

#first option
arr= np.where(df.eq('match'), df.columns, '').sum(axis=1)
arr=list(map(', '.join,arr))
print(arr)

#second option
arr= np.where(df.eq('match'), df.columns, '').sum(axis=1)
arr=list(map(list,arr))
print(arr)

Output:

#first option
['a, b', 'b', 'a, c']

#second option
[['a', 'b'], ['b'], ['a', 'c']]