根据包含列表中字符串的名称选择pandas数据框列

Question

我有一个数据框 df 和一个字符串列表 cols_needed，它们指示我想在 df 中保留的列。df 中的列名与 cols_needed 中的字符串并不完全匹配，因此我不能直接使用 intersection 等方法。但是，列名包含 cols_needed 中的字符串。我尝试了 str.contains，但无法使其正常工作。如何基于 cols_needed 对 df 进行子集划分？

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

- Gaurav Bansal

Answer 1

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

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

输出：

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

- Quang Hoang

Answer 2

通过filter的regex选项

df.filter(regex='|'.join(cols_needed))

---

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

- ALollz

Answer 3

列表推导也可以实现同样的功能：

columns = [cols for cols in df 
           for col in cols_needed 
           if col in cols]

['sim-prod1', 'sim-prod2']

In [110]: df.loc[:, columns]
Out[110]: 
   sim-prod1  sim-prod2
0          1          3
1          2          4

- sammywemmy

---

很好，或者在这种情况下只需使用 df[columns]。 - tdy