この記事は Pythonのコードを短く簡潔に書くテクニック Advent Calendar 2017 の4日目です。
はじめに
dictで一つのキーに複数の値を持つために値をリストにすることがあります。
この時、データを挿入していくには既に値があるかチェックし、なければ空のリストを挿入する、という操作が必要になります。
こういった操作をするのに便利なやり方をいくつか紹介します。
例題
単語が入ったリストから、頭文字をキーとしたdictに格納する処理を考えてみます。
元データの単語リスト
words = ['apple', 'orange', 'banana', 'alpha', 'beta']
↓
欲しい結果
results = {
'a': ['apple', 'alpha'],
'b': ['banana', 'beta'],
'o': ['orange'],
}
普通に書いた場合
results = {}
for word in words:
key = word[0]
if key not in results: #1
results[key] = [] #2
results[key].append(word) #3
普通に書くと、1.キーの存在チェック、2.空リストの作成、3.リストに単語を格納、の３回もdictにアクセスしなければなりません。
コードも長くなります。
dict.setdefault()を使う
setdefault()はdictから値を取得する際に指定のキーがなければデフォルト値を設定してくれます。
results = {}
for word in words:
results.setdefault(word[0], []).append(word)
これを使うとコードがスッキリする上に、dictへのアクセスが1回で済むようになります。
collections.defaultdictを使う
defaultdictでも同じようなことができます。
コンストラクタでデフォルト値を作成する関数を指定します。
[]でキーを参照したときに、なければデフォルト値が設定されます。
from collections import defaultdict
results = defaultdict(list)
for word in words:
results[word[0]].append(word)
itertools.groupbyを使う
groupbyを使うとリストを任意のキーでグルーピングできます。
この結果で辞書を作成すれば目的の結果が得られます。
ただしgroupbyに渡すリストは事前にソートする必要があります。
from itertools import groupby
from operator import itemgetter
results = {k: list(v)
for k, v in groupby(sorted(words), key=itemgetter(0))}
参考
- Python3 ドキュメント
- Python標準ライブラリ
- Qiita