Zenn
🎉

pythonのリスト内包表記における二重ループ

に公開
Python
内包表記
tech

はじめに

pythonでの内包表記は、もちろん多重ループも可能ですが、実行の度に迷うこともあるので整理がてらにまとめます。

内包表記とは

ほとんどの方が問題ないかと思いますが、ほんの少しばかり序章として書いておきます。

内包表記とは一言でいうと「forループでリストなどをシンプルに生成するための構文」といったところでしょうか。
例えば、0~4の整数のリストを作りたい場合

list = []
for i in range(5):
    list.append(i)

とすれば生成できます。これを内包表記で書くなら,

list = [i for i in range(5)]

と非常にシンプルに書くことができます。

内包表記の多重ループ

タイトルは多重ループとしましたが、ここでは二重ループで扱いたいと思います。

シンプルな二重ループ

内包表記の二重ループの例を1つ見てみます。

result = [(x, y) for x in [1, 2] for y in ['a', 'b']]

これは以下のような出力となります。

[(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]

つまり, 内包表記の内側のループ（変数xのほう）が固定→外側のループ（変数yのほう）が回るイメージということがわかりますね。
そのため、forループで書くならば

result = []
for x in [1, 2]:
    for y in ['a', 'b']:
        result.append((x, y))

ということになります。

応用例

ここからは実際の作業で必要になったところ。ネストの深いjsonデータなどと向き合うときに使えるかもしれません。

例としてこのようなデータを見ていきます。

  • data1：リストのデータ, 要素は辞書形式でのデータ, 長さは2
  • key_indices：リストのデータ, リストの中にリストで長さは2
data1 = [
    {
        'key1': 'v1a',
        'key2': 'v2a',
        'key3': [{'val': 'a0'}, {'val': 'a1'}, {'val': 'a2'}]
    },
    {
        'key1': 'v1b',
        'key2': 'v2b',
        'key3': [{'val': 'b0'}, {'val': 'b1'}, {'val': 'b2'}]
    }
]

key_indices = [[0, 1], [1, 2]]

ここでやりたいことは, data1から次のように取ってくることです。

[  
  [{'val': 'a0'}, {'val': 'a1'}],  
  [{'val': 'b1'}, {'val': 'b2'}]   
]

結論からいくと、以下の内包表記により取得ができます。

result = [
    [item['key3'][i] for i in indices]
    for item, indices in zip(data1, key_indices)
]

順を追って確認していきましょう。

  1. for item,indices in zip(data1,key_indices)について

つまり,

  • 1回目のループ；item←data1[0], indices←[0,1]
  • 2回目のループ；item←data1[1], indices←[1,2]
    が代入されます。
  1. [item['key3'][i]for i in indices]について

例えば, 上記の手順1での1回目のループの場合,変数iは[0,1]でループされるので

  • 1回目のループ；item['key3'][0]を取得
  • 2回目のループ；item['key3'][1]を取得
  1. 外側の内包表記について

結果として

[  
  [{'val': 'a0'}, {'val': 'a1'}],  
  [{'val': 'b1'}, {'val': 'b2'}]   
]

が得られるわけです。

まとめ

見た目が複雑な内包表記ですが、理解できるとすっきりしますね。
そしてコードもシンプルになるので使いこなしていきたいですね！

Discussion