pythonでList内のSetを特定の条件で高速に結合したい

Question

前提・実現したいこと

python v3.7.4で、List内のSetを特定の条件で高速に結合したいです。

具体的には、以下の様なListをインプットとして、
例) all_groups: List[Set[int]] = [{1, 2}, {2, 3}, {4, 5, 6}, {7, 8, 9}, {6, 8, 9}, {8, 10}, {11, 12}, {1, 2}, {2, 3}]

List内のSet同士を比較し、Set内に同一の数字が含まれる場合はSetを結合していき、最終的に
例) unified_groups: List[Set[int]] = [{1, 2, 3}, {4, 5, 6, 7, 8, 9, 10}, {11, 12}]
の様なリストを作成したいです。

実際には、all_groupsのサイズが数百万〜数千万件あり、メモリやCPUなどのリソースは潤沢という前提で、これをなるべく早く処理したいのですが、冴えた方法が思いつかず、良い方法があれば教えていただけると非常に助かります。

試したこと

昨日から、以下のコードで実行してみてるのですが、数千万件を処理すると1日経っても返ってこない状況です。

from typing import List, Set

def unify_group(all_groups: List[Set[int]]) -> List[Set[int]]:
  unified_groups: List[Set[int]] = []
  for group in all_groups:
    # 交差集合のリストを取得
    detected_unified_groups: List[Set[int]] = []
    for unified_group in unified_groups:
        if unified_group.isdisjoint(group) is False:
          detected_unified_groups.append(unified_group)

    if len(detected_unified_groups) == 0:
      # 交差集合がなければ、groupをそのまま追加
      unified_groups.append(group)
    else:
      # 交差集合があれば、既存グループを削除して、groupとdetected_unified_groupの和集合を追加
      unified_group: Set[int] = set()
      for detected_unified_group in detected_unified_groups:
        unified_groups.remove(detected_unified_group)
        unified_group = unified_group.union(detected_unified_group)
      unified_groups.append(unified_group)

  return unified_groups

Comments

[LouiS0616]

要素は必ず整数値ですか？また、上限・下限は既知ですか？

[rinoguchi]

確認ありがとうございます！
要素は必ず整数値で、上限・加減は既知です。
（下限は1で、上限は9000万ぐらいの値になります）

Answer 1

union-findというデータ構造を使う案はどうでしょうか。
実装はココの物を借りました。

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        return {r: self.members(r) for r in self.roots()}

    def __str__(self):
        return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())


all_groups = [{1, 2}, {2, 3}, {4, 5, 6}, {7, 8, 9}, {6, 8, 9}, {8, 10}, {11, 12}, {1, 2}, {2, 3}]
uf = UnionFind(13) # all_groupsに登場する最大値+1

for a, *others in all_groups:
    for b in others:
        uf.union(a, b)
print(uf)
"""
0: [0]
1: [1, 2, 3]
8: [4, 5, 6, 7, 8, 9, 10]
11: [11, 12]
"""

Comments

[rinoguchi]

ご回答ありがとうございます！
Union Find（素集合データ構造）と言うものがあるんですね。まさに今回のデータ構造そのものでした。
実際のデータで試して見たところ、unionするところまではすぐに終わるのですが、結果をList[Set[int]]の形で取得する処理を
`unified_groups: List[Set[int]] = list(map(lambda r: set(uf.members(r)), uf.roots()))`
こんな感じで実装したのですが、この部分にかなり時間がかかってしまいました。
お忙しいところご検討いただき、ありがとうございました。

Answer 2

無駄があるけと・・・

from collections import defaultdict

all_groups = [{1, 2}, {2, 3}, {4, 5, 6}, {7, 8, 9}, {6, 8, 9}, {8, 10}, {11, 12}, {1, 2}, {2, 3}]

d = {}
for v in all_groups:
    for k in v:
        if k in d:
            s = d[k]
            s.update(v)
            for c in v:
                if c in d and s != d[c]:
                    s.update(d[c])
            break
    else:
        s = set(v)
    for k in s:
        d[k] = s

print(list(map(set, set(map(tuple, d.values())))))

Comments

[rinoguchi]

ご回答ありがとうございます！
実データで試して見たところ、100万件=>100.7秒、200万件=>350.9秒で
自分の元の処理よりは圧倒的に速いものの、件数が増えると厳しくなりそうな感じでした。
お忙しいところ検討いただき、ありがとうございました。

Answer 3

(追記)
もっと素直に組んでみました

def unify_group2(all_groups: List[Set[int]]) -> List[Set[int]]:
    val_to_index = {}
    result = []

    for s in all_groups:
        indices = {val_to_index.get(n, -1) for n in s}
        if max(indices) == -1:
            result.append(s)
            index = len(result) - 1
            for n in s:
                val_to_index[n] = index
        else:
            indices = sorted(i for i in indices if i > -1)
            index = indices[0]
            for i in indices[1:]:
                result[index] |= result[i]
                for j in result[i]:
                    val_to_index[j] = index
                # 空をセットしておいて後で取り除く
                result[i] = {}
            result[index] |= s
            for n in s:
                val_to_index[n] = index

    return list(filter(None, result))

---

私だったら……

from typing import Dict, List, Set


def f(x: List[Set[int]]) -> List[Set[int]]:
    result: List[Set[int]] = []
    value_to_index: Dict[int, int] = {}  # 数値→resultのインデクス

    for s in x:
        for v in s:
            # 既出の要素であれば value_to_index に登録されている
            if v in value_to_index:
                idx = value_to_index[v]

                # 集合のマージ
                result[idx] |= s

                # 辞書に登録
                for vv in s:
                    value_to_index[vv] = idx
                break
        else:
            # ループが回りきった = 未知の要素しかない

            # sのコピーをresultに追加
            result.append(set(s))

            # 辞書に登録
            idx = len(result) - 1
            for v in s:
                value_to_index[v] = idx
    return result


def unify_group(all_groups: List[Set[int]]) -> List[Set[int]]:
    prev = []
    result = all_groups
    while len(result) != len(prev):
        prev = result
        result = f(prev)

    return result

既知の要素が出現したらマージする**だけ(↓)**のfをまず書いて、結果が収束するまで何度も何度も適用します。
(fは[{1,2},{3},{2,3}]を[{1,2,3},{3}]にするだけ)

確認用コード

def check(x):
    from itertools import permutations
    for s1, s2 in permutations(x, 2):
        assert not (s1 & s2)


def data():
    from random import randint, shuffle
    x = [set([randint(1, 1000) for j in range(randint(50, 80))]) for i in range(20000)]
    x += [set([randint(1000, 100_000) for j in range(randint(50, 80))]) for i in range(20000)]
    x += [set([randint(100_000, 200_000) for j in range(randint(50, 80))]) for i in range(2000)]
    x += [set([randint(200_001, 1_000_000) for j in range(randint(50, 80))]) for i in range(200)]
    shuffle(x)

    return x


print(
    unify_group(
        [{1, 2}, {2, 3}, {4, 5, 6}, {7, 8, 9}, {6, 8, 9}, {8, 10}, {11, 12}, {1, 2}, {2, 3}]
    )
)
r = unify_group(data())
print(len(r))
check(r)

Comments

[rinoguchi]

短期間に二つも回答いただきありがとうございます！
詳しくコメントまでかいていただき、とても分かりやすかったです。
実データで検証して見たところ、
一つ目：100万件=>10.4秒、200万件=>13.2秒、1000万件=>39.2秒
二つ目：100万件=>1.1秒、200万件=>2.4秒、1000万件=>13.1秒
と尋常じゃない速度で処理完了し、結果も全く問題ありませんでした。

いただいた回答の中で最も速く、しかも件数が増加しても時間は線形に増えるだけなので、数千万件の本番データを扱う場合でも全然大丈夫そうです。ベストアンサーにさせていただきました。