辞書を使ってLeetCodeの Two Sum問題をPythonで解く

はじめに

こんにちは、りょまです。プログラミング未経験ですが、エンジニアを目指してLeetCodeの問題を1日1題解き、その学びを記事にしています。

今回は「Two Sum」という問題に挑戦しました。
LeetCode: Two Sum

間違いや改善点があればご指摘いただけると幸いです。

問題の概要

与えられた配列 nums と整数 target に対して、nums の中から2つの数を足して target になる組み合わせを見つける問題です。

• 各入力には必ず1つの解が存在すると仮定してよいです
• 同じ要素を2回使うことはできません
• 答えはどの順序で返してもかまいません

例1

入力: nums: [1, 5, 3, 8, 7], target: 10
出力: [1, 4]（nums[1] + nums[4] = 5 + 5 = 10）

例2

入力: nums: [2, 9, 4, 6, 3], target: 12
出力: [1, 3]（nums[1] + nums[3] = 9 + 3 = 12）

制約

• 2 <= nums.length <= 10^4
• -10^9 <= nums[i] <= 10^9
• -10^9 <= target <= 10^9
• 1つの有効な答えが存在する

解法

辞書を利用して効率的に問題を解きます。（ハッシュマップとも言います。）

辞書とは、キーと値をセットで格納・取得できるデータ構造のことです。（英語で dictionary なので dict とします。）

# 辞書の例
dict = {'apple': 100, 'orange': 150, 'banana': 200}

# 値の取得
print(dict['apple']) # 出力結果: 100

# 値の追加
dict['lemon'] = 250

# 検索
print('lemon' in dict) # 出力結果: True

このようにしてデータの格納、値の取得や追加を実装することができます。上記の例では、キーにフルーツ名、値にそのフルーツの価格を格納しています。

例えば、appleの場合、キーが 'apple' で、値が 100 となります。

辞書の特徴は、3つあります。

• キーと値のペアを効率的に格納・取得できる
• キーは一意であり、重複することはない
• 高速な検索、挿入、削除操作ができる（平均的な時間計算量は O(1)）

具体的な解法

1. numsをループで1つずつ見ていく
2. targetからnumsの各数値を引いた値をcomplement(補数)として、そのcomplementが辞書にあるかを確認する
3. complementが存在する場合、そのインデックスと現在の数値のインデックスを返す
4. complementが存在しない場合、現在の数値とそのインデックスを辞書に追加する

辞書により、全ての要素を一度ずつ確認するだけで良いので、時間計算量はO(n)となる。
・forループの時間計算量: O(n)
・辞書操作の時間計算量: O(1)
・全体の時間計算量: O(n) × O(1) = O(n)

実際のコード

def two_sum(nums, target):
    dict = {}
    
    for i, num in enumerate(nums):
        complement = target - num
        if complement in dict:
            return [dict[complement], i]
        dict[num] = i

nums = [1, 5, 3, 8, 7]
target = 10
print(two_sum(nums, target))

nums = [2, 9, 4, 6, 3]
target = 12
print(two_sum(nums, target))

実行結果

[1, 4]
[1, 3]

コード説明

1. 辞書 dict を作成する
2. 配列 nums をループし、各要素に以下の操作を行う
   • target から現在の数値を計算し、complement に格納
   • if 文で complement が dict に存在するか確認
   • 存在する場合、complement のインデックスと現在のインデックスを返す
   • 存在しない場合、現在の数値とインデックスを辞書に追加する
3. forループが終了し、if文に引っ掛からなければ最終的に空の [] を返す

あまり馴染みのないコードがありました。

for i, num in enumerate(nums):

これはenumerate関数というもので、インデックス番号と任意の要素を同時に取り出すことができます。

subjects = ['Math', 'English', 'Science', 'Music']

for i, subject in enumerate(subjects):
    print(i, subject)

0 Math
1 English
2 Science
3 Music

このようにenumerate関数を使うことで配列subjectsに格納されているそれぞれの要素をインデックスと同時に取り出すことができました。

まとめ

今回の問題は辞書を使うことで、効率的に要素を見つけることができました。
辞書について理解が深まりとても勉強になりました。