アルゴリズム1000本ノック #7. 2 Sum

Problem

本問も LeetCode から拝借しています.

与えられた整数配列Aと1つの整数Kに対して, Aの中からその合計値がKとなるような
2つの整数を見つけ, それが先頭から何番目かを小さい順に配列の形で返却せよ.
（条件を満たす2つの整数は, 必ず配列A内に1組だけ存在するものと考えて良い）

入力: numbers={2, 7, 11, 15}, target=9
出力: [1, 2]

Solution

すべての組み合わせを総当りでチェックしていくと計算量 O(N^2) を要するので, 効率的な解き方が必要.

配列 A を先頭から順番に見ていき, 値とそれが先頭から何番目であるかをハッシュマップに格納していきます. こうすることで, 今見ている値に対して K - A[i] を満たすような値が既にあったかを, 同じハッシュマップから検索 することでO(1)で検出可能になるところがキーポイントですね.

Python版コード:

class Solution(object):
    def twoSum(self, nums, target):
        m = {} # store existence of each character
        for i, num in enumerate(nums):
            v = target - num
            if v not in m:
                m[num] = i + 1
            else:
                return [m[v], i+1]

計算量は O(N), 空間量も O(N) となります( N は配列Aの長さ）.

Java版コード:

public class Solution {
    public int[] twoSum(int[] nums, int target) {
        Map<Integer, Integer> m = new HashMap<Integer, Integer>();
        for (int i = 0; i < nums.length; i++) {
            int num = nums[i];
            int diff = target - num;
            if (m.containsKey(diff)) {
                return new int[] {m.get(diff), i+1};
            }
            m.put(num, i+1);
        }
        return new int[2];
    }
}