概要

Pythonにて二つのリスト同士を比較する際、リストの共通要素だけでなく、どちらのリストに何個共通要素が含まれているかを知りたい時があります。今回は、collectionsモジュールを使って、リストが長い場合に比較的高速で処理できる方法をまとめておきます。

背景

（読み飛ばしても大丈夫）なぜこんな処理をしたかったかというと、pandasのDataFrame同士をmergeしようとするときに、諸事情により双方の連携用のキーカラムが不明で、やむを得ず双方の全カラム同士をフルアタックでmergeして試さなければいけない、という酷い状況に直面したからです。
DataFrame内のカラムが少なかったり、レコードが短かったりすれば、すべてのカラムの組み合わせについて力業でmergeのチェックをすることもできるのですが、100列近いカラムが最大10万行以上のレコードを持っていますので、全てmergeしてチェックするのは現実的ではありません。しかも、組み合わせるDataFrameの数が100近くありました。サラリーマンとして「このデータセットでは仕事ができない」とは言えない状況に置かれていた関係上、劇的に効率的なチェック方法が必要でした。そのため、DataFrameからカラムだけをlistとして抽出して軽量なチェック処理ができないか、と考えた次第です。

困りごとの具体例

問題が無い場合の事例

リストが短い場合は実用上の問題は全くありません。

list_a = [1, 2, 3, 3, 4, 5]
list_b = [3, 4, 4, 6, 8, 9]

という二つのリストがあるとすると、両者に共通する要素は、

list_com = list(set(list_a)&set(list_b))
print(list_com)

を実行することで、

[3, 4]

という結果が即座に得られます。setでリストを集合にし、&で共通要素を抽出し、最後にリスト化する、というポピュラーな方法です。ただし、[3, 4]の2個の要素は、あくまで共通要素です。今回は、この共通要素がlist_aに含まれる個数を求めたいのですが、目視で分かるように[3, 3, 4]の３個ですから、目的とは異なります。
さて、Pythonでlist_aに含まれる要素数を数える場合、下記のようにリスト内包表記の長さを求めるのが最もシンプルな方法の一つでしょうか。

%time list_a_com = [x for x in list_a if x in list_com]
print(list_a_com)
print(f'len of common in list_a = {len(list_a_com)}')

その結果、下記の出力が得られます。

CPU times: total: 0 ns
Wall time: 0 ns
[3, 3, 4]
len of common in list_a = 3

上記の%timeはJupyter Labで実行時間を計測できるマジックコマンドです。%time以降のlist_a_com = [x for x in list_a if x in list_com]の処理時間を測定しているわけですが、ほぼ0秒であることがお分かりいただけると思います。

問題がある場合の事例

上記の短いリストの場合では処理時間がほぼ0sでしたが、リストが長い場合はそうはいきません。下記は、list_a、list_bとして、0から100000までの範囲の乱数100000個からなるリストを生成しています。

import random
list_a = [random.randint(0, 100000) for _ in range(100000)]
list_b = [random.randint(0, 100000) for _ in range(100000)]

リストが長すぎるので、中身ではなく長さだけ確認します。

print(f'len list_a = {len(list_a)}')
print(f'len list_b = {len(list_b)}')

その結果、

len list_a = 100000
len list_b = 100000

となります。最初のリストに比べると、4桁ほど長くなっています。
ここで、前と同様にlist_aとlist_b共通する要素を抽出します。ここも要素の個数だけ確認します。

%time list_com = list(set(list_a)&set(list_b))
print(f'len list_com = {len(list_com)}')

その結果、下記の出力が得られます。共通の要素を抽出するだけなら、十数ミリ秒と、ほとんど時間がかからないと分かります。

CPU times: total: 0 ns
Wall time: 17 ms
len list_com = 39849

続けて、list_aに含まれる共通要素の個数を調べます。

%time list_a_com = [x for x in list_a if x in list_com]
print(f'len={len(list_a_com)}')

その結果、下記のように、約30秒と、恐ろしく処理に時間がかかることになってしまいました（Core-i7のそれなりのスペックのマシンを使ってます）。

CPU times: total: 30.6 s
Wall time: 33.3 s
len=63015

一回の処理だけであれば深呼吸している間に終わってしまうので、気にしない、という手もありますね。ですが、自動で複数の組み合わせをチェックしたいといった場合、待ち時間がとんでもないことになってしまいます。実際に私は困る場面がありましたので、下記の対策を講じました。

解決策

冒頭で触れたとおり、collectionsモジュールを使って処理します。こちらのページが端的にまとまっていると思います。
まずはモジュールをimportします。

from collections import Counter

要素の出現回数を数えるための辞書型のサブクラス「Counter」です。

%time count = Counter(list_a)
print(count)

すると、10ミリ秒以下で出現個数の辞書が生成されます。例えば、先頭の59058:8は、list_aには59058が8個あった、という意です。

CPU times: total: 0 ns
Wall time: 9.03 ms
Counter({59058: 8, 30279: 7, 86287: 7, 37112: 7, 90718: 7, .....

この辞書から、共通する要素が何個list_aに含まれるかを調べます。具体的には、list_comの共通する要素のリストについて、順番にCounterで生成した辞書countからlist_aの重複した個数を読み出してリスト化し、最後に和を取っています。

%time sum([count[x] for x in list_com])

すると何と、10ミリ秒程度で要素数を求めることができました。

CPU times: total: 0 ns
Wall time: 11 ms
63015

まとめ

ちょっと工夫しただけでざっと1/1000の処理時間ですね。いろいろ調べて試して、という地道なアプローチが大事であることを実感した次第です。