コンサルでデータサイエンティスト

仕事でPythonを書いてます。機械学習、Webマーケティングに興味があります。趣味は旅です。

リスト内包表記を使ってPythonのリストの要素を条件付きで検索・抽出

コードを書く仕事をしていると、Pythonでリストの中から条件に一致する要素だけを抽出するようなスクリプトを書く機会も少なくないと思います。このようなとき、普通はforループを使って書いてしまいたくなるのですが、Pythonリスト内包表記を使えば、「これ1行で書けるんじゃね?」となります。for文との違いを示しながら実装例をご紹介します。

目次

Pythonのリストから要素を条件付きで抽出

次のような文字列のリストがあるとします。このリストから文字列などの条件をもとに特定の要素のみを抜き出したいとします。

txt_files = ['データサイエンス集中講義_20180912.txt', ¥
                  'データサイエンス集中講義_20180914.txt', ¥
                  'Pythonプログラミング_20180912.txt', ¥
                  'Pythonプログラミング_20180915.txt', ¥
                  'Pythonプログラミング_20180916.txt']

文字列が完全一致する要素を抽出

まずは 'データサイエンス集中講義_20180912.txt' と完全一致する要素を抽出し、新たなリスト output_list に入れたいとします。

forループで記述

普通に書くと次のようになります。

output_list = []

for file in txt_files:
    if file == 'データサイエンス集中講義_20180912.txt':
        output_list.append(file)
リスト内包表記で記述

リスト内包表記を使えば1行で次のように書くことができます。簡単ですよね?

[file for file in txt_files if 'データサイエンス集中講義_20180912.txt' == file]
コードの比較

各変数がどのように対応しているかを色別でみていくと理解が深まると思います。

forループ
> for file in txt_files:
> ____if file == 'データサイエンス集中講義_20180912.txt':
>________output_list.append(file)

リスト内包表記
> [file for file in txt_files if 'データサイエンス集中講義_20180912.txt' == file]


ある文字列を含む要素を抽出

今度は、'データサイエンス' という文字列を含む要素を抽出したいとします。in 句を使うと、特定の文字列が含まれているかどうかを分岐させることができます。

forループで記述

さきほどと同じように、forループを使って書くと次のようになります。

output_list = []
for file in txt_files:
    if 'データサイエンス' in file:
        output_list.append(file)
リスト内包表記で記述

リスト内包表記を使えば、こちらも1行ですっきりと書くことができます。

[file for file in txt_files if 'データサイエンス' in file]

文字列の一部を抽出

最後に、実務で発生しそうな例を紹介します。こちらではテキストファイルのファイル名のうち、日付部分のみを取り出して条件付きで抽出したい場合を考えます。

forループで記述

pythonは文字列変換に強いので、1つの文だけで多くのことができます。

output_list = []
for file in txt_files:
    if file.replace('.txt', '').split('_')[1] == '20180912':
        output_list.append(file)
リスト内包表記で記述

難しい処理も次のように1行で書くことができます。

[file for file in txt_files if '20180912' == file.replace('.txt', '').split('_')[1]]

まとめ

リスト内包表記を使ってPythonのリスト中身を条件付きで検索および抽出する方法についてご紹介しました。ある程度複雑な処理でも、リスト内包表記を使えば1行で書くことができることをお伝えしました。実は、リスト内包表記は1行で書けること以外にも処理速度が速いというメリットがあります。Pythonを書く機会が多い方は是非ともおさえておきましょう。