【LeetCode】1. Two Sumを解く

問題へのリンク

問題概要

整数からなる配列numsとある整数targetが与えられる．
このとき，足し合わせるとtargetになる２つの数のインデクスを返せ．

ただし，各入力はただ一つの解が存在し，同じ要素を2回使うことは出来ないものとする．

例：

Input: nums = [2,7,11,15], target = 9
Output: [0,1]
Output: Because nums[0] + nums[1] == 9, we return [0, 1]

制約

• $2 <= nums.length <= 103$
• $-109 <= nums[i] <= 109$
• $-109 <= target <= 109$
• ただ一つの解が存在する

考えたこと

まず最初に思いついたのは，2つのインデクスi,jを用意して，$nums[i]+nums[j]=target$となるような[i,j]を総当たりで求める方法だ：

class Solution(object):
    def twoSum(self, nums, target):
        """
        :type nums: List[int]
        :type target: int
        :rtype: List[int]
        """
        n = len(nums)
        for i in range(n):
            for j in range(i+1, n):
                if nums[i] + nums[j] == target:
                    return [i,j]

ただし，この解法は時間計算量が$O(n^2)$となり大きな問題の場合は計算時間超過となってしまう．（今回の制約では間に合う）

従って，何かしらの工夫をして計算量を減らす必要が出てくる．

$O(n)$の計算量で解を得ることを考えると，

1. numsの要素を前から順にみていく
2. $target - nums[i]$となるような要素がnumsに存在するか確かめる
3. 存在する場合：そのインデクスを返す
4. 存在しない場合：次の要素を探索

のように実装すればよさそうだ．

辞書型のdictに，入力numsに現れる整数nのインデクスindexを保存すると，

dict = {n: index, ...}

のようなハッシュテーブルが得られる．

このデータを利用することで，$target - nums[i]$が存在するか確かめる以下の操作

target - nums[i] in dict

が$O(1)$で行われる．

また，そのような値が存在したときはdict[target-nums[i]]によって定数時間でそのインデクスを得ることができる．

このような考えで実装した以下のコードは$O(n)$で実行される：

class Solution(object):
    def twoSum(self, nums, target):
        """
        :type nums: List[int]
        :type target: int
        :rtype: List[int]
        """
        dict = defaultdict(int)
        for i in range(len(nums)):
            if target - nums[i] in dict:
                return [i, dict[target - nums[i]]]
            dict[nums[i]] = i