Pythonでペアの合計をチェックするデータ構造を作成する方法

はじめに

ここでは、次の2つのメソッドを持つデータ構造を作成することを考えてみましょう。

• add(val): 値 val をデータ構造に追加します
• find(val): 合計が val となる2つの要素が存在するかどうかをチェックします

このデータ構造は、クエリが来たときに即座に結果を返せるように設計する必要があります。find() が呼ばれるたびにすべての数値の組み合わせを検索するような非効率な処理は避けたいところです。

例として、オブジェクト obj を作成し、6, 14, 3, 8, 11, 15 の数値を追加した後、obj.find(9)、obj.find(11)、obj.find(15) を呼び出すケースを考えます。このときの出力は True、True、False となります。9 は 6+3 で、11 は 3+8 で構成できるためです。一方、15 自体はデータ構造に存在しますが、合計が 15 になる2つの数の組み合わせは存在しません。

解決のアプローチ

この問題を解決するために、以下の手順に従います。

• コンストラクタを定義します。
• nums := 追加された数値を管理する新しいセット（集合）
• multiple := 重複して追加された数値を管理する新しいセット（集合）
• add() 関数を定義します。引数として val を受け取ります。
  • val がすでに nums に存在する場合は、val を multiple に挿入します。
  • それ以外の場合は、val を nums に挿入します。
• find() 関数を定義します。引数として val を受け取ります。
• nums 内の各要素 n に対して以下を繰り返します。
  • n + n が val と等しい場合、n が multiple に存在すれば True を返します（同じ数を2回使うケースへの対応）。
  • それ以外の場合、val - n が nums に存在すれば True を返します。
• ループが完了してもペアが見つからなければ False を返します。

実装例

より理解を深めるために、以下の実装を見てみましょう。

class PairSumChecker:
    def __init__(self):
        self.nums = set()
        self.multiple = set()

    def add(self, val):
        if val in self.nums:
            self.multiple.add(val)
        else:
            self.nums.add(val)

    def find(self, val):
        for n in self.nums:
            if n + n == val:
                return n in self.multiple
            elif val - n in self.nums:
                return True
        return False

obj = PairSumChecker()
obj.add(6)
obj.add(14)
obj.add(3)
obj.add(8)
obj.add(11)
obj.add(15)

print(obj.find(9))
print(obj.find(11))
print(obj.find(15))

入力

print(obj.find(9))
print(obj.find(11))
print(obj.find(15))

出力

True
True
False

仕組みのポイント

この実装の鍵となるのは、2つのセットを使い分ける点です。nums にはこれまでに追加された一意な数値が格納され、multiple には2回以上追加された数値が記録されます。これにより、「同じ数値を2回足す」ケース（例：7+7=14）でも、その数値が実際に2回追加されているかどうかを正確に判定できます。

計算量の面では、add() はセットへの挿入のみを行うため O(1) で動作し、find() は格納された数値を一度走査するため O(n) となります。事前準備のコストを最小限に抑えながら、クエリに対して高速に応答できるデータ構造が実現できるのです。