Python で groupby してみる (itertools のやつ)

f:id:monozukuri-bu:20171021104752j:plain

はじめまして、モノづくり部の古田です。
Python(Flask) での Web サービス開発に携わっています。

はじめに

groupby とは↓のことです。
https://docs.python.jp/3/library/itertools.html#itertools.groupby

今まで、いまいち使いどころが分からなかった itertools.groupby なんですが、 つい最近役に立つ場面があったので、ちょっと書いてみようと思います。

「もっと他に、こういう使い方があるよー」などの意見もあると思いますので、 そういった方は、お気軽にコメント頂けると有り難いです。m(_ _)m

利用場面

とても単純ですが「SQL が使えない状況で、集計的なことをしたい」場面です。
なので、SQL が使える状況で無理に使う必要は全くないですし、そんなことすると
「なんでわざわざ Python 使ってやってんだよ! カスが!!」
のような罵声を浴びるコトになりますのでご注意。

前提とサンプルデータ

チーム対抗でゲームを行い、チーム毎の獲得賞金をメンバで分配したいとします。
分配する際の基準は、メンバ毎の貢献ポイントです。
各メンバが、チーム内で何割貢献し、結果としていくら賞金がもらえるかを一覧で出力したいです。

・チーム毎の獲得賞金

チーム名 賞金
山田チーム 2,500,000
田中チーム 400,000
鈴木チーム 3,200,000

・メンバ毎の貢献ポイント

チーム名 メンバ名 貢献ポイント
山田チーム 山田 10
山田チーム 吉田 20
山田チーム 武田 30
田中チーム 田中 5
田中チーム 田辺 2
鈴木チーム 鈴木 80

・最終的に欲しいデータ

チーム名 賞金 メンバ名 貢献ポイント 貢献割合 取り分
山田チーム 2,500,000 山田 10 16.67% 416,750
山田チーム 2,500,000 吉田 20 33.33% 833,250
山田チーム 2,500,000 武田 30 50.00% 1,250,000
田中チーム 400,000 田中 5 71.43% 285,720
田中チーム 400,000 田辺 2 28.57% 114,280
鈴木チーム 3,200,000 鈴木 80 100.00% 3,200,000

コード全体

こんな感じです。

# groupby をインポート
from itertools import groupby

# 辞書型で定義
TEAM_PRIZE = {
    '山田チーム': 2500000,
    '田中チーム': 400000,
    '鈴木チーム': 3200000
}

# リスト型で定義
MEMBER_POINT = [
    ('山田チーム', '山田', 10),
    ('山田チーム', '吉田', 20),
    ('山田チーム', '武田', 30),
    ('田中チーム', '田中', 5),
    ('田中チーム', '田辺', 2),
    ('鈴木チーム', '鈴木', 80)
]

# 結果格納用
output_list = []

# groupby で回す
for key, group in groupby(sorted(MEMBER_POINT, key=lambda x: x[0]), key=lambda x: x[0]):

    # グループ化されたデータをリストに変換しておく
    group_list = list(group)

    # チーム内の貢献ポイント合計を出す
    team_point = sum(data[2] for data in group_list)

    # グループ化されたデータを for で回し、欲しいデータを output_list に入れていく
    for data in group_list:
        team_name = data[0]
        team_prize = TEAM_PRIZE[team_name]
        member_name = data[1]
        member_point = data[2]
        member_percentage = round(member_point * 100 / team_point, 2)
        member_prize = round(team_prize * member_percentage / 100)

        output_list.append((team_name, team_prize, member_name,
                            member_point, member_percentage, member_prize))

# 出力確認
[print(data) for data in output_list]

# 出力結果
# ('山田チーム', 2500000, '山田', 10, 16.67, 416750)
# ('山田チーム', 2500000, '吉田', 20, 33.33, 833250)
# ('山田チーム', 2500000, '武田', 30, 50.0, 1250000)
# ('田中チーム', 400000, '田中', 5, 71.43, 285720)
# ('田中チーム', 400000, '田辺', 2, 28.57, 114280)
# ('鈴木チーム', 3200000, '鈴木', 80, 100.0, 3200000)

補足

・グループ化されたデータの持ち方

以下の出力結果を見ると分かりやすいのですが、
key にはキーとした値、group にはキーによってグループ化された値がイテレータで格納されます。
そのため、二回以上読み込みたいときは、一度リストなどに変換しておくと楽です。

for key, group in groupby(MEMBER_POINT, key=lambda x: x[0]):
    print(key)
    print(group)
    print(list(group))

# 山田チーム
# <itertools._grouper object at 0x10d4790b8>
# [('山田チーム', '山田', 10), ('山田チーム', '吉田', 20), ('山田チーム', '武田', 30)]

# 田中チーム
# <itertools._grouper object at 0x10d479160>
# [('田中チーム', '田中', 5), ('田中チーム', '田辺', 2)]

# 鈴木チーム
# <itertools._grouper object at 0x10d4790b8>
# [('鈴木チーム', '鈴木', 80)]

・groupby する前にはソートが必要

公式ドキュメントにも書いてありますが、groupby は key の値を順次読んでいき、
それが変わったタイミングでグループを判断するため、key とする値でソートされていることが前提になります。