ハッシュマップでループ処理を高速化してみる　#450

leetcodeでハッシュマップ実装の練習をしたのでアウトプットします。
お題は以下です。

整数の配列 nums と整数のターゲットが与えられたとき、numsのうち2つの数値の合算がターゲットになるような組み合わせを返す。

numsは正確に1つの解を持つと仮定してよく、同じ要素を2度使ってはならない。

答えを返す順番は自由である。

leetcode

入力値と答えの例は以下です。

-- Example１--------
Input: nums = [2,7,11,15], target = 9
Output: [0,1]
Explanation: Because nums[0] + nums[1] == 9, we return [0, 1].
 
-- Example2--------
Input: nums = [3,2,4], target = 6
Output: [1,2]

まずはハッシュマップを使わずに、単純に全文検索して答えを探すコードを書いてみます。

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        for index_first in range(len(nums)):
            for index_second in range(len(nums)):
                if index_first == index_second:
                    continue
                if nums[index_first] + nums[index_second] == target:
                    return [index_first, index_second]

結果は以下で、かなり時間がかかっています。
Runtime: 3880 ms
Memory Usage: 17.4 MB

forループを2つ入れ子で回し、足し算をしてtarget数値になればそのインデックスを返す、という方法です。

つぎにハッシュマップを使ってみます。
各数値の「募集する数値（足すとtargetになる数値）」をキーに持ち、募集元である自身のインデックスを要素に持ちます。

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        nums_to_add = {}
        for index, num in enumerate(nums):
            num_to_add = target - num
            nums_to_add[num_to_add] = index

        for index, num in enumerate(nums):
            if num in nums_to_add:
                if index == nums_to_add[num]:
                    continue
                else:
                    return [index, nums_to_add[num]]

これだけで驚くほど高速化しました。
Runtime: 57 ms
Memory Usage: 18.1 MB

つぎにハッシュマップを更に効率化してみます。
先の例では一旦全ての数値に対してnums_to_addを作成していますが、以下のように都度判定しながら処理すると、最低限のループ数で処理可能です。

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        nums_to_add = {}
        for index, num in enumerate(nums):
            num_to_add = target - num
            if num in nums_to_add:
                return [index, nums_to_add[num]]
            nums_to_add[num_to_add] = index

結果は以下で、最も高速になりました。
Runtime: 47 ms
Memory Usage: 18 MB

ここではループの最初で、まずは自身が募集されている数字かどうかを判定します。

募集されていれば自身のインデックスと募集元のインデックスを返します。募集されていなければ、募集したい数字をキーに、募集元である自身のインデックスを要素に設定して次のループに入ります。

なぜ初めから思いつかなかったか不思議なくらいシンプルで、高速な処理になりました。

ここまでお読みいただきありがとうございました！