Pythonでリストの連続する同じ値の要素をグループ化(itertools.groupby)

Modified: | Tags: Python, リスト

Pythonで、リストなどのイテラブルオブジェクトの連続(隣接)する同じ値の要素をまとめてグループ化するには、標準ライブラリitertoolsのgroupby()を使う。

import itertools

l = [0, 0, 0, 1, 1, 2, 0, 0]
print([(k, list(g)) for k, g in itertools.groupby(l)])
# [(0, [0, 0, 0]), (1, [1, 1]), (2, [2]), (0, [0, 0])]

順番(連続する・しない)に関わらず、同じ値の要素の個数を数え上げるにはcollections.Counterを使う。以下の記事を参照。

itertools.groupby()の使い方

itertools.groupby()はキーとグループのイテレータを返す。そのままprint()しても値は表示されない。

l = [0, 0, 0, 1, 1, 2, 0, 0]
print(itertools.groupby(l))
# <itertools.groupby object at 0x110ab58b0>

返されるグループもイテレータ。例えばlist()などでリスト化して使う。

返されるグループはそれ自体がイテレータで、 groupby() と iterable を共有しています。もととなる iterable を共有しているため、 groupby() オブジェクトの要素取り出しを先に進めると、それ以前の要素であるグループは見えなくなってしまいます。従って、データが後で必要な場合にはリストの形で保存しておく必要があります: itertools.groupby() --- 効率的なループ実行のためのイテレータ生成関数 — Python 3.11.3 ドキュメント

for k, g in itertools.groupby(l):
    print(k, g)
# 0 <itertools._grouper object at 0x110a26940>
# 1 <itertools._grouper object at 0x110a2c400>
# 2 <itertools._grouper object at 0x110aa8f10>
# 0 <itertools._grouper object at 0x110aa8ee0>

for k, g in itertools.groupby(l):
    print(k, list(g))
# 0 [0, 0, 0]
# 1 [1, 1]
# 2 [2]
# 0 [0, 0]

キーのみ、グループのみ、両方(キーとグループのタプル)をリストとして取得するには、リスト内包表記を利用すればよい。

print([k for k, g in itertools.groupby(l)])
# [0, 1, 2, 0]

print([list(g) for k, g in itertools.groupby(l)])
# [[0, 0, 0], [1, 1], [2], [0, 0]]

print([(k, list(g)) for k, g in itertools.groupby(l)])
# [(0, [0, 0, 0]), (1, [1, 1]), (2, [2]), (0, [0, 0])]

値を判定する関数を指定: 引数key

itertools.groupby()には引数keyを指定できる。使い方はsorted()max(), min()など他の関数のkeyと同じ。

keyに指定した関数(呼び出し可能オブジェクト)の結果を元に、連続した要素の値が同じかどうかが判定される。

例えば、文字列のリストに対して文字列の長さ(文字数)を返す組み込み関数len()を指定すると、文字数が同じ要素がグループ化される。keyに指定する際はカッコ()は不要。

l = ['aaa', 'bbb', 'ccc', 'a', 'b', 'aa', 'bb']
print([(k, list(g)) for k, g in itertools.groupby(l, len)])
# [(3, ['aaa', 'bbb', 'ccc']), (1, ['a', 'b']), (2, ['aa', 'bb'])]

偶数か奇数かで判定する例。ラムダ式(無名関数)を使っている。

l = [0, 2, 0, 3, 1, 4, 4, 0]
print([(k, list(g)) for k, g in itertools.groupby(l, lambda x: x % 2)])
# [(0, [0, 2, 0]), (1, [3, 1]), (0, [4, 4, 0])]

SQLのGROUP BYのような集約処理

SQLにおけるGROUP BYのように二次元データ(リストのリストなど)のある列を基準にグループ化したい場合も引数keyを指定する。

ここではリストやタプルの任意の位置(n番目)の要素を取得するのにラムダ式を使っているが、operator.itemgetter()を使うこともできる。

また、出力が見やすいようにfor文を使っているが、もちろん、これまでの例のようにリスト内包表記を使ってもよい。

l = [[0, 'Alice', 0],
     [1, 'Alice', 10],
     [2, 'Bob', 20],
     [3, 'Bob', 30],
     [4, 'Alice', 40]]

for k, g in itertools.groupby(l, lambda x: x[1]):
    print(k, list(g))
# Alice [[0, 'Alice', 0], [1, 'Alice', 10]]
# Bob [[2, 'Bob', 20], [3, 'Bob', 30]]
# Alice [[4, 'Alice', 40]]

itertools.groupby()では連続する同じ値の要素のみをグループ化する。順番に関係なくグループ化するには、元のリストなどをsorted()でソートしておく。

リストのリストをソートする場合、デフォルトでは各リストの先頭の要素を基準にソートされる。任意の位置の要素でソートするにはsorted()の引数keyを指定する。

for k, g in itertools.groupby(sorted(l, key=lambda x: x[1]), lambda x: x[1]):
    print(k, list(g))
# Alice [[0, 'Alice', 0], [1, 'Alice', 10], [4, 'Alice', 40]]
# Bob [[2, 'Bob', 20], [3, 'Bob', 30]]

さらに数値などを集計する例は以下の通り。ジェネレータ式を用いている。

for k, g in itertools.groupby(sorted(l, key=lambda x: x[1]), lambda x: x[1]):
    print(k, sum(x[2] for x in g))
# Alice 50
# Bob 50

なお、pandasにもグループ化・集約のためのgroupby()がある。複雑なデータ処理はpandasのほうが便利。

リスト以外(タプルや文字列など)の場合

リストだけでなく、タプルや文字列なども同様にitertools.groupby()で処理が可能。

タプルの場合。

t = (0, 0, 0, 1, 1, 2, 0, 0)
print([(k, list(g)) for k, g in itertools.groupby(t)])
# [(0, [0, 0, 0]), (1, [1, 1]), (2, [2]), (0, [0, 0])]

グループをリストではなくタプルにしたい場合はtuple()、全体をタプルにしたい場合はジェネレータ式とtuple()を使う。

print(tuple((k, tuple(g)) for k, g in itertools.groupby(t)))
# ((0, (0, 0, 0)), (1, (1, 1)), (2, (2,)), (0, (0, 0)))

文字列の場合。

s = 'aaabbcaa'
print([(k, list(g)) for k, g in itertools.groupby(s)])
# [('a', ['a', 'a', 'a']), ('b', ['b', 'b']), ('c', ['c']), ('a', ['a', 'a'])]

グループを文字列にしたい場合はjoin()を使う。

print([(k, ''.join(g)) for k, g in itertools.groupby(s)])
# [('a', 'aaa'), ('b', 'bb'), ('c', 'c'), ('a', 'aa')]

もちろん、イテラブル(反復可能)オブジェクトであればこれらの例以外も同様に処理できる。

関連カテゴリー

関連記事