Two Sum問題とは?Pythonを用いて解説

プログラミング

はじめに

プログラミングを学ぶ際の定番問題であるTwo Sum問題。これはリスト内から2つの数値を見つけ、その合計が特定の目標値と一致するようにする課題だ。

この問題を解くことで、データ構造やアルゴリズムの基礎を深く理解することができる。この記事では、効率的なPythonコードを例に、この問題の解法とポイントを説明する。


問題の概要

Two Sum問題は以下を満たすインデックスのペアを探すものである。

  • 与えられたリストnums内の2つの異なる要素の合計がtargetに一致する。
  • 見つかった場合、その2つのインデックスを返す。

例えば、nums = [2, 7, 11, 15]target = 9の場合、出力は[0, 1]となる。これは、nums[0] + nums[1] = 9を満たすためだ。


実装コード

以下はPythonによる効率的な実装である。

def two_sum(nums, target):
    # 数値とそのインデックスを格納する辞書
    num_dict = {}

    # リストをループして各数値を確認
    for i, num in enumerate(nums):
        complement = target - num

        # 辞書に補数が存在する場合、そのインデックスと現在のインデックスを返す
        if complement in num_dict:
            return [num_dict[complement], i]

        # 現在の数値を辞書に登録
        num_dict[num] = i

    # 解が見つからない場合
    return None

# 使用例
nums = [2, 7, 11, 15]
target = 9
result = two_sum(nums, target)
print(result)  # 出力: [0, 1]

コードの解説

1. 辞書num_dictの活用

num_dictは数値をキー、インデックスを値として保持する。これにより、特定の数値が存在するかを素早く確認できる。辞書の検索は平均してO(1)の計算量を持つため、高速である。

2. 補数の計算

complement = target - num

現在の要素numと目標値targetの差を計算することで、もう1つ必要な数値(補数)を特定する。

3. 補数の確認とインデックスの返却

if complement in num_dict:
    return [num_dict[complement], i]

補数がすでに辞書に存在する場合、インデックスのペアを即座に返却する。この仕組みにより、二重ループを回避できる。

4. 現在の要素を辞書に登録

num_dict[num] = i

次の探索に備えて、現在の数値とそのインデックスを辞書に追加する。


実際の動作シミュレーション

以下に動作を示す。

イテレーション 現在の数値 (num) 補数 (complement) 辞書 (num_dict) 結果
1 2 7 {} 辞書に登録
2 7 2 {2: 0} ペア発見 [0,1]

計算量の分析

  • 時間計算量
    リストを1度だけループするため、計算量はO(n)。これは入力サイズnに対して線形であり、非常に効率的なアルゴリズムである。

  • 空間計算量
    使用する辞書のサイズがリストの長さに比例するため、空間計算量もO(n)。


Two Sum問題の解法バリエーション

Two Sum問題は、そのシンプルさゆえに多様な解法が存在する。それぞれの手法には特徴があり、用途や制約に応じて使い分けが可能だ。以下に代表的な解法をいくつか紹介する。


1. 二重ループによる全探索

最も直感的な方法は、全てのペアを確認することだ。
この方法は単純だが、効率は良くない。

実装

def two_sum_brute_force(nums, target):
    for i in range(len(nums)):
        for j in range(i + 1, len(nums)):
            if nums[i] + nums[j] == target:
                return [i, j]
    return None

特徴

  • 時間計算量: O(n²)
    全てのペアを試すため、要素数が多いと非効率。
  • 空間計算量: O(1)
    辞書などの追加データ構造を使用しない。

2. 辞書を使った線形探索(ハッシュマップ方式)

効率的な解法として広く使われる方法。これが前述のコードのベースとなる。

特徴

  • 時間計算量: O(n)
    辞書を使い、補数の存在をO(1)で確認する。
  • 空間計算量: O(n)
    数値とインデックスを保存するための辞書が必要。

3. ソートと二分探索を組み合わせた方法

リストをソートし、二分探索を活用する方法。

実装

def two_sum_with_binary_search(nums, target):
    nums_sorted = sorted((num, i) for i, num in enumerate(nums))  # 数値と元のインデックスを保持
    for i in range(len(nums_sorted)):
        complement = target - nums_sorted[i][0]
        left, right = i + 1, len(nums_sorted) - 1

        # 二分探索で補数を探す
        while left <= right:
            mid = (left + right) // 2
            if nums_sorted[mid][0] == complement:
                return [nums_sorted[i][1], nums_sorted[mid][1]]
            elif nums_sorted[mid][0] < complement:
                left = mid + 1
            else:
                right = mid - 1
    return None

特徴

  • 時間計算量: O(n log n)
    ソートにO(n log n)、二分探索にO(log n)。
  • 空間計算量: O(n)
    ソート結果を保持するための追加メモリが必要。

注意

この方法では元のインデックスを追跡するため、ソート前のインデックスも保持する工夫が必要。


4. 双方向ポインタを使った探索

リストをソートし、左右のポインタを使って探索を進める方法。

実装

def two_sum_with_two_pointers(nums, target):
    nums_sorted = sorted((num, i) for i, num in enumerate(nums))
    left, right = 0, len(nums_sorted) - 1

    while left < right:
        current_sum = nums_sorted[left][0] + nums_sorted[right][0]
        if current_sum == target:
            return [nums_sorted[left][1], nums_sorted[right][1]]
        elif current_sum < target:
            left += 1
        else:
            right -= 1
    return None

特徴

  • 時間計算量: O(n log n)
    ソートにO(n log n)、ポインタの操作はO(n)。
  • 空間計算量: O(n)
    ソート結果を保持する必要がある。

5. 制約に応じた解法の調整

問題の制約が変わる場合、以下のような調整が必要となる。

(1) 重複要素を許容しない場合

リスト内で同じ要素が複数回使われることを防ぐ必要がある。

def two_sum_no_duplicates(nums, target):
    num_dict = {}
    for i, num in enumerate(nums):
        if target - num in num_dict and num_dict[target - num] != i:
            return [num_dict[target - num], i]
        num_dict[num] = i
    return None

(2) 全てのペアを見つける場合

通常のTwo Sumは最初の解を見つけた時点で終了するが、全てのペアを探す場合には以下のようにする。

def two_sum_all_pairs(nums, target):
    num_dict = {}
    results = []
    for i, num in enumerate(nums):
        complement = target - num
        if complement in num_dict:
            results.append((num_dict[complement], i))
        num_dict[num] = i
    return results

6. カスタマイズ例。インデックスではなく値を返す

特定のケースでは、インデックスではなく値そのものを返すことが求められる場合がある。

実装

def two_sum_return_values(nums, target):
    num_dict = {}
    for num in nums:
        complement = target - num
        if complement in num_dict:
            return [complement, num]
        num_dict[num] = True
    return None

解法の選択基準

  1. 入力サイズが小さい場合

    • 全探索でも十分処理可能。
  2. 効率を優先する場合

    • 辞書を使った線形探索。
  3. インデックスではなく値や全ペアが必要な場合

    • 問題の要件に応じてカスタマイズ。

Two Sumの応用例

Two Sum問題は単純なアルゴリズム問題の一つだが、その基本的な考え方を応用することで、より複雑な課題に対応できる。以下に、Two Sumの考え方を応用した代表的な問題例を挙げる。


1. Three Sum問題

問題概要

Three Sumは、リスト内の3つの異なる数値の和が特定の値に一致する組み合わせを見つける問題だ。例えば、nums = [-1, 0, 1, 2, -1, -4]target = 0 の場合、出力は [-1, -1, 2][-1, 0, 1] などになる。

解法

Three Sumでは、Two Sumを内部で利用することで効率的に解ける。

def three_sum(nums, target):
    nums.sort()
    results = []
    for i in range(len(nums) - 2):
        if i > 0 and nums[i] == nums[i - 1]:  # 重複を避ける
            continue
        left, right = i + 1, len(nums) - 1
        while left < right:
            current_sum = nums[i] + nums[left] + nums[right]
            if current_sum == target:
                results.append([nums[i], nums[left], nums[right]])
                while left < right and nums[left] == nums[left + 1]:
                    left += 1
                while left < right and nums[right] == nums[right - 1]:
                    right -= 1
                left += 1
                right -= 1
            elif current_sum < target:
                left += 1
            else:
                right -= 1
    return results

# 使用例
nums = [-1, 0, 1, 2, -1, -4]
target = 0
print(three_sum(nums, target))  # 出力: [[-1, -1, 2], [-1, 0, 1]]

2. 4 Sum問題

問題概要

リスト内の4つの異なる数値の和が特定の値に一致する組み合わせを見つける問題。Three Sumと似ているが、組み合わせの数が増えるため、効率的なアルゴリズムが求められる。

解法

Three Sumのアルゴリズムを拡張し、外側のループを1段増やすことで解く。

def four_sum(nums, target):
    nums.sort()
    results = []
    for i in range(len(nums) - 3):
        if i > 0 and nums[i] == nums[i - 1]:  # 重複を避ける
            continue
        for j in range(i + 1, len(nums) - 2):
            if j > i + 1 and nums[j] == nums[j - 1]:  # 重複を避ける
                continue
            left, right = j + 1, len(nums) - 1
            while left < right:
                current_sum = nums[i] + nums[j] + nums[left] + nums[right]
                if current_sum == target:
                    results.append([nums[i], nums[j], nums[left], nums[right]])
                    while left < right and nums[left] == nums[left + 1]:
                        left += 1
                    while left < right and nums[right] == nums[right - 1]:
                        right -= 1
                    left += 1
                    right -= 1
                elif current_sum < target:
                    left += 1
                else:
                    right -= 1
    return results

# 使用例
nums = [1, 0, -1, 0, -2, 2]
target = 0
print(four_sum(nums, target))  # 出力: [[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]

3. Two Sumのカウント問題

問題概要

Two Sumのペアを探すだけでなく、ペアの総数を数える問題。リストに重複要素が含まれる場合、同じペアを複数回数えることがあるため、効率的な実装が求められる。

解法

辞書を活用し、カウントを効率的に行う。

def two_sum_count(nums, target):
    num_count = {}
    count = 0
    for num in nums:
        complement = target - num
        if complement in num_count:
            count += num_count[complement]
        num_count[num] = num_count.get(num, 0) + 1
    return count

# 使用例
nums = [1, 1, 2, 2, 3]
target = 4
print(two_sum_count(nums, target))  # 出力: 2 (ペア: [1,3], [2,2])

4. K Sum問題

問題概要

リスト内でk個の数値の和が目標値になる組み合わせを求める問題。Three SumやFour Sumを一般化した形である。

解法

再帰を使い、kの値に応じて分岐する。以下はK Sumの一般的なフレームワーク。

def k_sum(nums, target, k):
    nums.sort()
    def find_k_sum(nums, target, k):
        if k == 2:
            results = []
            left, right = 0, len(nums) - 1
            while left < right:
                current_sum = nums[left] + nums[right]
                if current_sum == target:
                    results.append([nums[left], nums[right]])
                    left += 1
                    right -= 1
                elif current_sum < target:
                    left += 1
                else:
                    right -= 1
            return results
        results = []
        for i in range(len(nums) - k + 1):
            if i > 0 and nums[i] == nums[i - 1]:
                continue
            for subset in find_k_sum(nums[i + 1:], target - nums[i], k - 1):
                results.append([nums[i]] + subset)
        return results
    return find_k_sum(nums, target, k)

# 使用例
nums = [1, 0, -1, 0, -2, 2]
target = 0
k = 4
print(k_sum(nums, target, k))  # 出力: [[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]

5. Two Sumにおける制約追加の応用例

重複した組み合わせを許可しない場合

リスト内の数値を1回だけ使用する制約を課す場合、辞書を更新しないか、すでに使用した数値をスキップする処理を追加する。

動的配列でのリアルタイム更新

動的にリストの内容が変化する場合(例えば、数値がリアルタイムで追加される)、辞書やセットを使用して逐次的にTwo Sumの解を探索できる。


注意点とベストプラクティス

  1. 配列に重複がある場合
    このアルゴリズムは正しく動作する。辞書に登録する際にインデックスが上書きされるが、補数検索の仕組みがそれを補う。

  2. インデックスの順序
    辞書に登録された順序に基づき、返されるインデックスの順序が決まる。

  3. 解が存在しない場合
    上記コードでは、解が見つからない場合Noneを返す設計となっている。


まとめ

Two Sum問題は、辞書を活用することで効率的に解くことができる。計算量はO(n)と最適で、Pythonの基本構文を活用しつつ、データ構造の重要性を学ぶ良い機会となる。
この手法を理解することで、他のアルゴリズム問題にも応用できる知識が身につく。

コードを実際に書き、自分で試すことでさらなる理解が深まるはずだ。興味があれば、さらに複雑なバリエーション問題にも挑戦してみてほしい。

コメント