バイオインフォをやっている大学院生です。プログラミングのことなどを書きます。

itertoolsを使って列挙した場合に、要素のインデックスを取得する方法

itertoolsを使って組み合わせ列挙などを行う場合で、同時に取り出された要素のそれぞれのリスト内でのインデックスも取得したい場合の方法です。

まずenumerateを組み合わせてやるとこうなります。

import itertools

l1 = ["a", "b", "c"]
l2 = ["d", "e", "f"]

for (i1, v1), (i2,v2) in itertools.product(enumerate(l1), enumerate(l2)):
    print("index1:",i1, "value1:", v1, "index2:",i2, "value2:", v2)

# 出力
# index1: 0 value1: a index2: 0 value2: d
# index1: 0 value1: a index2: 1 value2: e
# index1: 0 value1: a index2: 2 value2: f
# index1: 1 value1: b index2: 0 value2: d
# index1: 1 value1: b index2: 1 value2: e
# index1: 1 value1: b index2: 2 value2: f
# index1: 2 value1: c index2: 0 value2: d
# index1: 2 value1: c index2: 1 value2: e
# index1: 2 value1: c index2: 2 value2: f

enumerateで生成したタプルを変数に展開するため、(i1,v1)のように囲む必要があるのに注意してください。

forループを下のように書いても同じようになります。インデックスの組み合わせを作成して、値を取得します。

for i1, i2 in itertools.product(range(len(l1)), range(len(l2))):
    v1 = l1[i1]
    v2 = l2[i2]
    print("index1:",i1, "value1:", v1, "index2:",i2, "value2:", v2)

また、加えてforループが何回回ったかを取得するなら次のようにかけます。展開する変数がごちゃごちゃで嫌なので、シンプルに書ける方法ご存知だったら教えてください…

for i, ((i1, v1), (i2,v2)) in enumerate(itertools.product(enumerate(l1), enumerate(l2))):
    print(i, "index1:",i1, "value1:",v1, "index2:", i2, "value2:", v2)

#  出力
# 0 index1: 0 value1: a index2: 0 value2: d
# 1 index1: 0 value1: a index2: 1 value2: e
# 2 index1: 0 value1: a index2: 2 value2: f
# 3 index1: 1 value1: b index2: 0 value2: d
# 4 index1: 1 value1: b index2: 1 value2: e
# 5 index1: 1 value1: b index2: 2 value2: f
# 6 index1: 2 value1: c index2: 0 value2: d
# 7 index1: 2 value1: c index2: 1 value2: e
# 8 index1: 2 value1: c index2: 2 value2: f