Pythonでzipを使ってJSONライクなdictをいい感じにforループで回す
Pythonでzip関数を使ってJSONライクな辞書（dict）を、いい感じにforループで要素ごとに抽出する例を解説します。ほぼワンライナーでできるので簡単です。
想定
こんなJSONライクな辞書があったとします。JSONからパースした場合でもいいですね。
data = {
"in_channels": [1024, 1024, 512, 256, 128],
"out_channels": [1024, 512, 256, 128, 64],
"upsample": [True, True, True, True, True],
"resolution": [8, 16, 32, 64, 128],
"attention": {
8: False, 16:False, 32:False, 64:True, 128:False
}
}
これをforループで回して、各配列の中身を取得したいのです。例えば、こんな出力にしたいのです。
(1024, 1024, True, 8, False)
(1024, 512, True, 16, False)
(512, 256, True, 32, False)
(256, 128, True, 64, True)
(128, 64, True, 128, False)
実はこれはほぼワンライナーでできます。それを見ていきましょう。
zip関数の基本
まずはzip関数の基本から。zip関数は
for a, b in zip([1, 2, 3], [4, 5, 6]):
print(a, b)
とすると
1 4
2 5
3 6
のように出力するのがzip関数です。要はzip関数の中に複数の一次元配列を放り込めればいいのです。
ただしdictの配列数（要素数）は可変なので、真面目にzip(a, b, …)のように書くのはアホらしいです。どうすればいいでしょう？
アンパックのアスタリスク「＊」
答えは、シーケンスのアンパック・アスタリスク（＊）を使いましょう。可変長引数でも使われるやつですね。zipと組み合わせると先程の例は次のようにも書けます。
data = [[1, 2, 3], [4, 5, 6]]
for a, b in zip(*data):
print(a, b)
これは先程と同じ出力になります。
1 4
2 5
3 6
dictの値だけ配列の配列に格納するには、dict.values()を使えばいいのでこれでできそうな気がします。
ネストされた辞書のケア
辞書の中身がすべて配列だったらOKですが、values()→アンパック→zipだとネストされた辞書が要素にあったときに困ります。例えば、
def main():
data = {
"in_channels": [1024, 1024, 512, 256, 128],
"out_channels": [1024, 512, 256, 128, 64],
"upsample": [True, True, True, True, True],
"resolution": [8, 16, 32, 64, 128],
"attention": {
8: False, 16:False, 32:False, 64:True, 128:False
}
}
for seq in zip(*data.values()):
print(seq)
このようにすると、
(1024, 1024, True, 8, 8)
(1024, 512, True, 16, 16)
(512, 256, True, 32, 32)
(256, 128, True, 64, 64)
(128, 64, True, 128, 128)
ネストされた辞書（attention）のキーだけ抽出されてしまいました。ネストされた辞書の値がほしいのでちょっとこれは困りました。
そこで以下のようにループ側で対策します。
# ネストされた辞書対策
def main():
data = {
"in_channels": [1024, 1024, 512, 256, 128],
"out_channels": [1024, 512, 256, 128, 64],
"upsample": [True, True, True, True, True],
"resolution": [8, 16, 32, 64, 128],
"attention": {
8: False, 16:False, 32:False, 64:True, 128:False
}
}
for seq in zip(*(v.values() if type(v) is dict else v for v in data.values())):
print(seq)
単純にvalues()で出てきた値がdictかどうかを調べているだけです。もしvがネストされた辞書だったら、さらにそれのvalues()を読みに行くことで、値側を取りに行くようにしています。
結果は以下の通りです。
(1024, 1024, True, 8, False)
(1024, 512, True, 16, False)
(512, 256, True, 32, False)
(256, 128, True, 64, True)
(128, 64, True, 128, False)
これで目的の結果になりました。
注意点
この例はPython3.7以上で行っています。Python3.7以上ではdictの順番が保証されるのが明確に言語仕様として追加されたのでOrderedDictを使わなくてもキーの順番保証ができるようになっています。Pythonのバージョンが古い場合は、OrderedDictを使って順番のケアをしてある必要があるかもしれません（あるいはライブラリによってはOrderedDictでパースしているケースがあるので、そのタイプのケアなど）。