LeetCode #1.Two SumをPythonで解いてみた

問題

Given an array of integers, return indices of the two numbers such that they add up to a specific target.
You may assume that each input would have exactly one solution, and you may not use the same element twice.

https://leetcode.com/problems/two-sum/

Example:

Given nums = [2, 7, 11, 15], target = 9,
Because nums[0] + nums[1] = 2 + 7 = 9,
return [0, 1].

整数のリストとターゲット数が与えられるので、整数のリストの中から足し合わせるとターゲット数になる2つの数を見つけ、そのインデックスを返せという問題です。答えは1つの組み合わせしか無く、同じインデックスの数２つを返すことは出来ません。

解法

1. ブルートフォース | 計算量O(n²)

ただ単純にforループを二回回すだけです。返却する答えは合っているはずですが、Time Limit Exceededを食らってLeetCode上では採点してもらえませんでした。

def twoSum(self, nums: List[int], target: int) -> List[int]:
    for i in range(len(nums)):
        for j in range(1, len(nums)):
            if (i != j) and (nums[i] + nums[j] == target):
                return [i, j]

テスト結果

2.HashMap利用 | 計算量O(n)

受け取ったリストを一周しつつ、key:数値 value:数値のインデックスのmapを作っていく方法です。今のポインタよりも前に答えがあればO(1)の計算量ですぐに答えを返却できるため、ブルートフォースよりも効率的です。
Python3を使った解法の上位5~6%くらいの速度であったようです。

def twoSum(self, nums: List[int], target: int) -> List[int]:
    hashmap = {}
    for i, num in enumerate(nums):
        search_key = target - num
        if search_key not in hashmap:
            hashmap[num] = i
        else:
            return [i, hashmap[search_key]]

テスト結果

注意点

mapにkey:valueの組み合わせを挿入する処理は、既にmap内に同じkeyが存在していないかチェックしてからやらなければなりません。存在チェックをせずにそのまま更新していくような処理にしてしまうと、数値リスト[3,3],target=6のようなテストケースで答えを返せなくなります。

数値リスト[3,3]、target=6で何も返却できずエラーになる例

def twoSum(self, nums: List[int], target: int) -> List[int]:
    hashmap = {}
    for i, num in enumerate(nums):
        hashmap[num] = i
        search_key = target - num
        if (search_key in hashmap) and (i != hashmap[search_key]):
            return [i, hashmap[search_key]]

まとめ

解説より

We reduce the look up time from O(n)O(n) to O(1)O(1) by trading space for speed. A hash table is built exactly for this purpose, it supports fast look up in near constant time.

https://leetcode.com/problems/two-sum/solution/

HashMapとはO(n)の計算量がかかる処理をO(1)で済むようにするための機構である。

感想

これまでHashMapを計算量を削減するためのものだと意識することなく使っていたため、このような見方を得ることが出来たのは収穫であった。