Python の集合を使ったアルゴリズム実践

集合の O(1) 検索を活かすと、多くのアルゴリズムを効率的に実装できる。ここでは実践的な例を紹介する。

2 つの配列の共通要素を見つける

リスト同士の共通要素を探す問題は、積集合で簡単に解ける。

def find_common(list1, list2):
    return list(set(list1) & set(list2))

a = [1, 2, 3, 4, 5]
b = [4, 5, 6, 7, 8]
print(find_common(a, b))  # [4, 5]

リストのままネストしたループで探すと O(n²) だが、set を使えば O(n) で済む。

配列内で和が特定の値になるペアを探す

Two Sum 問題は、見た要素を set で追跡することで効率化できる。

def two_sum(numbers, target):
    seen = set()
    pairs = []
    for num in numbers:
        complement = target - num
        if complement in seen:
            pairs.append((complement, num))
        seen.add(num)
    return pairs

nums = [1, 2, 3, 4, 5, 6]
print(two_sum(nums, 7))  # [(3, 4), (2, 5), (1, 6)]

各要素について「target から引いた値」が既出かどうかを O(1) で判定できる。

連続する整数の最長列を見つける

配列から最長の連続整数列を見つける問題がある。

def longest_consecutive(nums):
    num_set = set(nums)
    longest = 0
    
    for num in num_set:
        # 連続列の開始点のみ処理
        if num - 1 not in num_set:
            current = num
            length = 1
            
            while current + 1 in num_set:
                current += 1
                length += 1
            
            longest = max(longest, length)
    
    return longest

nums = [100, 4, 200, 1, 3, 2]
print(longest_consecutive(nums))  # 4（1, 2, 3, 4）

num - 1 が存在しない要素だけを起点にすることで、各要素を高々 2 回しか処理しない。

重複のない部分文字列の最長

スライディングウィンドウと set を組み合わせた古典的な問題だ。

def longest_unique_substring(s):
    seen = set()
    left = 0
    max_length = 0
    
    for right in range(len(s)):
        while s[right] in seen:
            seen.remove(s[left])
            left += 1
        seen.add(s[right])
        max_length = max(max_length, right - left + 1)
    
    return max_length

print(longest_unique_substring("abcabcbb"))  # 3（"abc"）
print(longest_unique_substring("bbbbb"))      # 1（"b"）
print(longest_unique_substring("pwwkew"))     # 3（"wke"）

ウィンドウ内の文字を set で管理し、重複が現れたらウィンドウを縮小する。

グラフの探索（BFS/DFS）

訪問済みノードの管理に set が最適だ。

def bfs(graph, start):
    visited = set()
    queue = [start]
    result = []
    
    while queue:
        node = queue.pop(0)
        if node not in visited:
            visited.add(node)
            result.append(node)
            queue.extend(graph.get(node, []))
    
    return result

graph = {
    'A': ['B', 'C'],
    'B': ['A', 'D', 'E'],
    'C': ['A', 'F'],
    'D': ['B'],
    'E': ['B', 'F'],
    'F': ['C', 'E']
}

print(bfs(graph, 'A'))  # ['A', 'B', 'C', 'D', 'E', 'F']

訪問済みチェックが O(1) なので、大規模グラフでも効率的に動作する。

集合の差分を使ったフィルタリング

許可リストや禁止リストによるフィルタリングは差集合で表現できる。

all_users = {"alice", "bob", "charlie", "dave", "eve"}
banned_users = {"charlie", "eve"}

active_users = all_users - banned_users
print(active_users)  # {'alice', 'bob', 'dave'}

複数の条件を組み合わせることも容易だ。

premium_users = {"alice", "dave", "frank"}
active_users = {"alice", "bob", "charlie", "dave"}

# プレミアムかつアクティブ
print(premium_users & active_users)  # {'alice', 'dave'}

# プレミアムまたはアクティブ
print(premium_users | active_users)  # {'alice', 'bob', 'charlie', 'dave', 'frank'}

アナグラム判定

2 つの文字列がアナグラムかどうかは、文字の出現を比較すればわかる。

from collections import Counter

def is_anagram(s1, s2):
    return Counter(s1) == Counter(s2)

print(is_anagram("listen", "silent"))  # True
print(is_anagram("hello", "world"))    # False

ただし、使われている文字の種類だけを見るなら set で十分だ。

def same_chars(s1, s2):
    return set(s1) == set(s2)

print(same_chars("abc", "cba"))    # True
print(same_chars("aab", "abb"))    # True（文字の種類は同じ）

素数の篩（エラトステネスの篩）

素数判定にも set が使える。

def sieve_of_eratosthenes(n):
    primes = set(range(2, n + 1))
    
    for i in range(2, int(n ** 0.5) + 1):
        if i in primes:
            # i の倍数を除去
            primes -= set(range(i * i, n + 1, i))
    
    return sorted(primes)

print(sieve_of_eratosthenes(30))
# [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]

差集合 -= で倍数を一括削除できるのが便利だ。

部分和問題

与えられた数の組み合わせで特定の合計を作れるか判定する問題がある。

def can_sum(target, numbers):
    reachable = {0}
    
    for num in numbers:
        new_reachable = set()
        for r in reachable:
            new_reachable.add(r + num)
        reachable |= new_reachable
    
    return target in reachable

print(can_sum(7, [2, 3]))     # True（2+2+3）
print(can_sum(7, [5, 3, 4]))  # True（3+4）
print(can_sum(7, [2, 4]))     # False

到達可能な合計を set で管理し、新しい数を加えるたびに更新する。

順序なしペアの生成

重複なく順序を無視したペアを生成するのに frozenset が使える。

items = ['a', 'b', 'c']
pairs = set()

for i in items:
    for j in items:
        if i != j:
            pairs.add(frozenset([i, j]))

print(pairs)
# {frozenset({'a', 'b'}), frozenset({'a', 'c'}), frozenset({'b', 'c'})}

frozenset を使うことで、(a, b) と (b, a) が同じペアとして扱われる。