CSVファイルにおける特定の列群だけ取得したい

Question

実現したいこと

・条件を満たす列の抽出

前提

環境はJupyterLabで、pythonを使用しています

csvファイルに大量の列があり、各列名は"ID+α"であり、αは列ごとに異なります

列の中から特定のIDを名前に含んだ列だけを抽出したファイルを作成したいです

特定のIDとは4桁の数で、それらは配列で与えられています

#列名の例
columns = "0001_A","0002_B","0003_C", ... ,"0027_AA", ...
#特定のID群の例
arr = ["0001","0002","0005","0012", ... ] 

いくつかの方法を調べましたが、上記を容易に解決できる方法が思い浮かびませんでした

発生している問題・エラーメッセージ

皆さんのお力をお借りできれば幸いです
よろしくおねがいいたします

調査したこと・試したこと

列名の一部にIDが含まれているため、pandasのfilterでlikeを行おうとしましたが
likeは文字列でしか使えないため、配列ごとは適用できませんでした

また、列を一つずつ抜き出したファイルを作成し、条件に沿ったファイルだけを結合しようとしましたが
複数ファイルの読み込みかつキー結合が必要なため、pandasや、スクリプトでも難しく断念いたしました

Answer 1

read_csv()でusecolを使って、csvを読み込むときに列を絞りこむ方法もあります。
列名を引数にとって使いたい列に対してはTrueを返す関数を渡します。

import pandas as pd
from io import StringIO

s = '''\
0001_A,0002_B,0003_C,0004_D,0005_E
1,2,3,4,5
6,7,8,9,10
'''

arr = ['0001', '0004']

df = pd.read_csv(StringIO(s), usecols=lambda col: any(col.startswith(i) for i in arr))

print(df)
#   0001_A  0004_D
#0       1       4
#1       6       9

Answer 2

pandasを使った例は出ているので、 csvを使ったやつを。

泥臭い感じで作ってますが、ポイントは
columns = [colname for colname in reader.fieldnames if colname.split('_')[0] in col_list]
こうやって、必要なカラムを抽出してるところですかね。

import csv

col_list = ["0001","0003","0012"] 

with open('test.csv', mode='r') as inf:
    reader = csv.DictReader(inf)
    columns = [colname for colname in reader.fieldnames if colname.split('_')[0] in col_list]

    with open('outf.csv', mode='w', newline='') as outf:
        writer = csv.DictWriter(outf, fieldnames=columns)
        writer.writeheader()
        for org_items in reader:
            items = {col:val for col, val in org_items.items() if col in columns}
            writer.writerow(items)```

Comments

[754x]

回答ありがとうございます。
こちらの方でも確認したところ、無事動きました。
今回は最も回答時間が早い方をベストアンサーとさせていただきました。

Answer 3

以下のような感じで欲しいID軍の列colsを抽出してやればよいと思います。

import pandas as pd

df = pd.DataFrame([[1,2,3]], columns = ['0001_A','0002_B','0003_C'])
print(df)
#   0001_A  0002_B  0003_C
#0       1       2       3

#特定のID群の例
arr = ['0001','0002','0005']

cols = list(filter(lambda c:list(filter(lambda s: s in c, arr)), df.columns))
ret = df[cols]
print(ret)
#   0001_A  0002_B
#0       1       2

Comments

[754x]

回答ありがとうございます。
こちらの方でも確認したところ、無事動きました。
今回は最も回答時間が早い方をベストアンサーとさせていただきました。

Answer 4

import pandas as pd

#列名の例
columns = ["0001_A","0002_B","0003_C","0004_D","0005_E","0006_F",
           "0007_G","0008_H","0009_I","0010_J","0011_K","0012_L"]
#特定のID群の例
arr = ["0001","0002","0005","0012"] 

#
df = pd.DataFrame(data=[[*range(len(columns))]], columns=columns)
dfx = df.filter(regex='|'.join(arr))
print(dfx)

#    0001_A  0002_B  0005_E  0012_L
# 0       0       1       4      11

Comments

[754x]

回答ありがとうございます。
こちらの方でも確認したところ、無事作動し、目的を達成できました。

最も回答時間が早かったためベストアンサーとさせていただきました。
ありがとうございました。