pandas фильтрует значения столбцов по нескольким значениям

Question

У меня есть df с несколькими столбцами, такими как MLB, NBA, NHL, NFL, TESTNBA. Я хотел бы вернуть список, в котором столбцы содержат строку MLB или NBA. так как показано ниже:

df_check = ['MLB', 'NBA', 'TESTNBA']

value_cols = [col for col in df.columns if df_check in col]

Вышеприведенный сбой со следующей ошибкой:

TypeError: для 'in' в качестве левого операнда требуется строка, а не список

Есть ли способ отфильтровать столбцы по нескольким значениям?

Answer 1

Вы можете использовать pandas.DataFrame.filter:

Настройка

df = pd.DataFrame(columns=['MLB', 'NBA', 'NHL', 'NFL', 'TESTNBA'])    
df_check = ['MLB', 'NBA']

---

df.filter(regex='|'.join(df_check)).columns

Index(['MLB', 'NBA', 'TESTNBA'], dtype='object')

Answer 2

Оператор in не работает в векторизованном виде. С вашей логикой вам нужно будет использовать any с выражением генератора:

value_cols = [col for col in df.columns if any(i in col for i in df_check)]

Лучше, вы можете использовать логическое индексирование с str.contains. Вот демо:

df = pd.DataFrame(columns=['MLB1', 'NBA2', 'SOMEOTHERCOL', 'TESTNBA3', 'MLB4'])

df_check = ['MLB', 'NBA', 'TESTNBA']

value_cols = df.columns[df.columns.str.contains('|'.join(df_check))]

Index(['MLB1', 'NBA2', 'TESTNBA3', 'MLB4'], dtype='object')