RSSフィードを購読
Python Coding Challenge - How will you find two numbers whose sum equals the target?
(Python コーディングチャレンジ [和が target に等しい２つの要素をより速い実行時間で見つけよう！] 編)
(Python コーディングチャレンジ [和が target に等しい２つの要素をより速い実行時間で見つけよう！] 編)
シンプルな実装方法では O(n²) になりますが dictionary を上手く使用することで O(n) で同じ処理をこなせます。分かりますか？
問題：
int 値を要素とする list と、int 値 target が与えられます。この時、足し算の結果が target と等しくなる list 内の２つの要素を見つけ、それらの要素のインデックス番号からなる list を返すプログラムを記述してください。
与えられる list 内に含まれる「答え」は常に「１つだけ」存在します。また、同じ要素を２回使ってはいけません。
例：
与えられるリスト nums: [3, 6, 7, 10, 11]
target: 14
target: 14
# 出力: [0, 4]
解説：
nums[0] + nums[4] == 3 + 11 == 14 == target です。よってそれぞれの要素のインデックス番号からなるリスト [0, 4] を返します。7 + 7 も 14 ですが同じ要素を２回使うことはできません。
また for ループのネストを利用する以下の実装は認められません。この実装は問題なく動作しますが実行計算時間 (time complexity) は O(n²) です:
def two_sum(nums: list[int], target: int) -> list[int]:
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
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) でこの問題を解決するアルゴリズムを実装してください。そのためには dictionary を上手く利用する必要があります。
また関数定義は以下のような形です (同じである必要はありません):
def two_sum(nums: list[int], target: int) -> list[int]:
...
...
作成したプログラムは以下のテストをパスする必要があります:
import unittest
class TestTwoSum(unittest.TestCase):
def test_1(self):
nums = [3, 8, 13, 15]
target = 11
res = two_sum(nums, target)
self.assertEqual(res, [0, 1])
def test_2(self):
nums = [3, 8, 13, 15]
target = 16
res = two_sum(nums, target)
self.assertEqual(res, [0, 2])
def test_3(self):
nums = [3, 8, 13, 15]
target = 18
res = two_sum(nums, target)
self.assertEqual(res, [0, 3])
def test_4(self):
nums = [3, 8, 13, 15]
target = 21
res = two_sum(nums, target)
self.assertEqual(res, [1, 2])
def test_5(self):
nums = [3, 8, 13, 15]
target = 23
res = two_sum(nums, target)
self.assertEqual(res, [1, 3])
def test_6(self):
nums = [3, 8, 13, 15]
target = 28
res = two_sum(nums, target)
self.assertEqual(res, [2, 3])
def test_7(self):
nums = [3, 6, 7, 10, 11]
target = 14
res = two_sum(nums, target)
self.assertEqual(res, [0, 4])
class TestTwoSum(unittest.TestCase):
def test_1(self):
nums = [3, 8, 13, 15]
target = 11
res = two_sum(nums, target)
self.assertEqual(res, [0, 1])
def test_2(self):
nums = [3, 8, 13, 15]
target = 16
res = two_sum(nums, target)
self.assertEqual(res, [0, 2])
def test_3(self):
nums = [3, 8, 13, 15]
target = 18
res = two_sum(nums, target)
self.assertEqual(res, [0, 3])
def test_4(self):
nums = [3, 8, 13, 15]
target = 21
res = two_sum(nums, target)
self.assertEqual(res, [1, 2])
def test_5(self):
nums = [3, 8, 13, 15]
target = 23
res = two_sum(nums, target)
self.assertEqual(res, [1, 3])
def test_6(self):
nums = [3, 8, 13, 15]
target = 28
res = two_sum(nums, target)
self.assertEqual(res, [2, 3])
def test_7(self):
nums = [3, 6, 7, 10, 11]
target = 14
res = two_sum(nums, target)
self.assertEqual(res, [0, 4])
考え方:
問題に掲載した for ループのネストを利用するアルゴリズムにおいて、外側のループ i は与えられたリストの最初から最後までを辿っていきます。そして内側のループ j で i 以降の要素を最後まで辿っています。結局 i と j の実行計算時間はそれぞれ O(n) ですからそれが２重となっていることでプログラム全体の実行計算時間が O(n²) になっているわけです。
ここで何とか内側のループをなくして実行計算時間を O(1) にすることができれば、プログラム全体の time complexity が O(n) になります。
もう一度このプログラムで何をやっているのかを考えてみましょう。ループ i で取り出した要素に「加える」と target になる数、つまり target に対する nums[i] の補数 (complement) を探しているわけです。そして for ループのネストを使っている方法は、ループ i で見つけた数に対する補数を「それ以降」の要素の中から選び出す、ということをやっているわけですね。
ここで考え方を 180° 変えて、ループ i で見つけた数に対する補数を「それまでに出てきた」要素の中から選び出す、としたらどうでしょう？
まず最初に空の dictionary; nums_dict を定義しておきます。ループ i で要素を取り出すたびに target に対する nums[i] の補数が nums_dict に保存されているか、をチェックします。もし存在していなければ nums[i]: i という key: value ペアで nums_dict に保存し次のループへと進みます。もし存在していれば「答え」が見つかったわけですから、お互いのインデックス番号からなるリストを作成、返してプログラムは終了します。
問題に掲載した例題で考えてみましょう。nums_dict を作成します; {}。ループ i は nums; [3, 6, 7, 10, 11] の各要素を辿っていきます。target は 14 です:
i -> 0, nums[i] -> 3, 補数は 14 - 3 -> 11, nums_dict は空ですから 11 は存在しません。保存します; nums_dict -> {3: 0}
i -> 1, nums[i] -> 6, 補数は 14 - 6 -> 8, nums_dict には 8 は存在しません。保存します; nums_dict -> {3: 0, 6: 1}
i -> 2, nums[i] -> 7, 補数は 14 - 7 -> 7, nums_dict には 7 は存在しません。保存します; nums_dict -> {3: 0, 6: 1, 7: 2}
i -> 3, nums[i] -> 10, 補数は 14 - 10 -> 4, nums_dict には 4 は存在しません。保存します; nums_dict -> {3: 0, 6: 1, 7: 2, 10: 3}
i -> 4, nums[i] -> 11, 補数は 14 - 11 -> 3, nums_dict に 3 が存在します。ここで「答え」が見つかりました。[nums_dict[3], i] -> [0, 4] を返してプログラムは終了です。
プログラム実装例：
def two_sum(nums: list[int], target: int) -> list[int]:
nums_dict = {}
for i in range(len(nums)):
if (comp := target - nums[i]) in nums_dict:
return [nums_dict[comp], i]
nums_dict[nums[i]] = i
return None
nums_dict = {}
for i in range(len(nums)):
if (comp := target - nums[i]) in nums_dict:
return [nums_dict[comp], i]
nums_dict[nums[i]] = i
return None
Complexity Analysis
(アルゴリズム解析)
Time complexity (実行計算時間):
Dictionary への値の挿入、値の検索、値の取得の実行時間は全て O(1) です。つまりこのアルゴリズムを採用することで外側のループの内側の操作を O(n) から O(1) にすることに成功したわけです。よってプログラム全体の time complexity も O(n) になります。
Space complexity (実行メモリ要求量):
最悪で要素数 n 分の dictionary を作成する必要がありますからその分のメモリが必要となります。よって space complexity は O(n) となります。つまり、for ループのネストを利用するアルゴリズムと dictionary を利用するアルゴリズムでは「時間」と「空間」のトレードオフを行なっていることになります。