Pythonで2要素のリストを共通要素ごとにグループ化する3つの方法

各サブリストが2つの要素を持つ「リストのリスト」を想定します。サブリストの片方の要素は、他の多くのサブリストと共通しています。この記事では、共通要素ごとにサブリストをグループ化した最終的なリストを作成する方法を、3つのアプローチで解説します。

方法1: setとmapを使う

サンプルのリストでは、最初の要素が文字列（曜日）、2番目の要素が数値になっています。まずmapとlambdaで各サブリストの2番目の要素を抽出し、setに変換することで重複のないキーの一覧を作成します。その後、各キーに対して元のリストを走査し、2番目の要素が一致するサブリストから1番目の要素を取り出してグループ化します。

コード例

listA = [['Mon', 2], ['Tue', 3], ['Wed', 3],
         ["Thu", 1], ['Fri', 2], ['Sat', 3],
         ['Sun', 1]]

# setとmapを使用
temp = set(map(lambda i: i[1], listA))
res = [[j[0] for j in listA if j[1] == i] for i in temp]

# 結果の出力
print("グループ化された要素のリスト : \n", res)

出力結果

上記のコードを実行すると、次の結果が得られます。

グループ化された要素のリスト :
[['Thu', 'Sun'], ['Mon', 'Fri'], ['Tue', 'Wed', 'Sat']]

方法2: groupbyとitemgetterを使う

operatorモジュールのitemgetter関数で各サブリストの2番目の要素を取得し、その結果をキーとしてitertools.groupbyを適用します。groupbyは隣接する同一キーの要素だけをまとめる仕組みのため、事前にsortで同じキー順に並べ替えておく必要がある点に注意してください。

コード例

from itertools import groupby
from operator import itemgetter

listA = [['Mon', 2], ['Tue', 3], ['Wed', 3],
         ["Thu", 1], ['Fri', 2], ['Sat', 3],
         ['Sun', 1]]

# groupbyを使用（事前にソートが必要）
listA.sort(key=itemgetter(1))
groups = groupby(listA, itemgetter(1))
res = [[i[0] for i in val] for (key, val) in groups]

# 結果の出力
print("グループ化された要素のリスト : \n", res)

出力結果

上記のコードを実行すると、次の結果が得られます。

グループ化された要素のリスト :
[['Thu', 'Sun'], ['Mon', 'Fri'], ['Tue', 'Wed', 'Sat']]

方法3: defaultdictを使う

collectionsモジュールのdefaultdictを使うと、サブリストの2番目の要素を辞書のキーとして扱えます。通常のdictと違い、存在しないキーにアクセスしてもKeyErrorが発生せず、自動的に空のリストが生成されるため、1番目の要素をそのままappendできるのが特徴です。コードが最もシンプルになります。

コード例

import collections

listA = [['Mon', 2], ['Tue', 3], ['Wed', 3],
         ["Thu", 1], ['Fri', 2], ['Sat', 3],
         ['Sun', 1]]

# defaultdictを使用
res = collections.defaultdict(list)
for val in listA:
    res[val[1]].append(val[0])

# 結果の出力
print("グループ化された要素のリスト : \n", res)

出力結果

上記のコードを実行すると、次の結果が得られます。

グループ化された要素のリスト :
defaultdict(<class 'list'>, {2: ['Mon', 'Fri'], 3: ['Tue', 'Wed', 'Sat'], 1: ['Thu', 'Sun']})

出力はdefaultdictオブジェクトですが、dict(res)のように変換すれば、通常の辞書として扱うことも可能です。

まとめ

いずれの方法でも同じグループ化結果が得られます。簡潔に書きたい場合はsetとmapの組み合わせ、大量のデータを効率的に処理したい場合は1回のループで完結するdefaultdictがおすすめです。groupbyを使う場合は、事前のソートが必須である点に注意しましょう。