Zenn
🐧

#1:LeetCode-TwoSumをやってみる

に公開
Python
LeetCode
twosum
tech

問題(翻訳)

整数配列 nums と整数 target が与えられています。
合計が target になるような 2 つの数のインデックス を返してください。
各入力には 必ず 1 つの解 が存在すると仮定してよいものとします。
同じ要素を 2 回使用することはできません。
答えは 任意の順序 で返して構いません。

制約

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

解法

初めに思いついた解法

  • 方針
    • ブルートフォース(ナイーブ解法)で打開する方法
two_sum.py
class Solution:
    def twoSum(self, 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 []

  • 時間計算量・空間計算量

    • 時間的計算量：for文を2重で実行し、配列を走査しているためにO(n^2)となる
    • 空間的計算量：入力の値のみを使用した走査のためO(1)となる

  • メモ

    • 今回のような制約がある場合2重のループ((10^4)^2=10^8≈100000000)だとTLEする可能性がある
    • しかし、使用しているメモリについては定数倍のため効率は良く見える
    • 必ず有効回答が存在するが、一応空のリストを返すようにしておく

最終的な解法

  • 方針
    • hashmapを使用し、該当する値とインデックスをペアで記憶しておく方法
two_sum.py
class Solution:
    def twoSum(self, nums: list[int], target: int) -> list[int]:
        num_to_idx = {}
        for idx, num in enumerate(nums):
            if num in num_to_idx:
                return [num_to_idx[num], idx]
    
            num_to_idx[target - num] = idx
            
        return []

  • 時間計算量・空間計算量

    • 時間的計算量：走査が1重のforになり、O(n)に改善されている
    • 空間的計算量：num_to_idxが用意されたことで、入力に比例する値を格納する可能性があるためO(n)となる

  • メモ

    • 時間的計算量は改善されたが、その分空間的計算量が増えた

テストケース

s = Solution()

test_cases = [
    (1, [2,7,11,15], 9, [0, 1]),
    (2, [3,2,4], 6, [1, 2]),
    (3, [3,3], 6, [0, 1]),
]

for no, nums, target, expected in test_cases:
    got = s.twoSum(nums, target)
    assert (
        got == expected
    ), f"Test{no}, nums: {nums}, target: {target}, got: {got}, expected: {expected}"

学び・気づき

  • 今回はdictを使用することで、時間計算量をO(n)に減少できた反面、空間的計算量がO(n)に増加した。計算量のトレードオフにはもう少し理解が必要
  • 制約を意識することでロジックの改良に目がいくようになった
  • アルゴリズムも大事だけど、データ構造をどこでどう使うかも大事

参考

GitHubで編集を提案

Discussion