itertools.groupbyを使ったグルーピング・ブレイク処理

何かと気づきが得られるモジュールitertoolsですが、あるリストを走査しながら特定の条件でブレイクして出力を行うと行った処理を上手く記述する方法を検討します。

例えば次のようなリスト(deptでソート済み)があり、dept別にファイルに出力するケースの場合、どのように行うか考えてみます。

emp_list = [
  {'dept': 10, 'ename': 'scott'},
  {'dept': 10, 'ename': 'john'},
  {'dept': 10, 'ename': 'tom'},
  {'dept': 20, 'ename': 'bob'},
  {'dept': 20, 'ename': 'maria'},
]

簡単に思いつくのは、順番に走査しながらdeptが変わったタイミングでファイルに出力する or ファイルを切り替えるといったものです。
以下のコードは、emp_listを上から走査しながら、deptが変わったタイミングでファイルに出力。変わらなかったら貯めこむという例です。

# deptが変わるまで溜め込んでから、ファイルに書き出す例
out = []
for emp in emp_list:
  if len(out) > 0 and out[0]['dept'] != emp['dept']:
    # deptがブレイクした時
    with open("dept-%d.txt" % out[0]['dept'], "wt") as f:
      f.write("\n".join(str(e) for e in out))
    out.clear()
  out.append(emp)
if len(out) != 0:
  # 残ったoutを出力
  with open("dept-%d.txt" % out[0]['dept'], "wt") as f:
    f.write("\n".join(str(e) for e in out))

見た感じ分かりにくいコードですが、この方式のイマイチな点は、ブレイク条件「deptが変わる」でファイル出力する時と、ブレイクしないまま最後に残ったものをファイルに出力する時のコードが重複している点です。groupbyを使うともっと綺麗に書けます。

from itertools import groupby

dept_emp_groups = groupby(emp_list, lambda e:e['dept'])
for dept, dept_group in dept_emp_groups:
    with open("dept-%d.txt" % dept, "wt") as f:
        f.write("\n".join(str(e) for e in dept_group))

groupbyを使うと指定したリスト(listでなくてイテレーション可能なもの)から、特定のキー別にグルーピングした結果を得ることができます。
(実際にはlistではなくイテレータですが)こんなイメージです。

[
  (10, [{'dept': 10, 'ename': 'scott'},{'dept': 10, 'ename': 'john'},  {'dept': 10, 'ename': 'tom'}],
  (20, [{'dept': 20, 'ename': 'bob'},  {'dept': 20, 'ename': 'maria'}],
]

deptをキーにグルーピングしているため、存在するdept別に処理を行います。deptが0件であっても問題ありません。

参考:10.1. itertools — 効率的なループ実行のためのイテレータ生成関数 — Python 3.5.3 ドキュメント

関連記事:

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

日本語が含まれない投稿は無視されますのでご注意ください。(スパム対策)