python - 根据列表中包含字符串的名称选择熊猫数据框列

Question

我有一个数据框df和一个字符串列表cols_needed，它们表示我想要保留的列df。中的列名df与 中的字符串不完全匹配cols_needed，因此我不能直接使用类似intersection. 但列名确实包含cols_needed. 我试着玩弄，str.contains但无法让它工作。我如何df根据cols_needed?

import pandas as pd
df = pd.DataFrame({
    'sim-prod1': [1,2],
    'sim-prod2': [3,4],
    'sim-prod3': [5,6],
    'sim_prod4': [7,8]
})

cols_needed = ['prod1', 'prod2']

# What I want to obtain:
    sim-prod1  sim-prod2
0      1        3
1      2        4

标签： pythonpandas

Answer 1

您可以使用联合模式进行探索str.contains，例如：

df.loc[:,df.columns.str.contains('|'.join(cols_needed))]

输出：

   sim-prod1  sim-prod2
0          1          3
1          2          4