two sum (hash で O(n))｜コードで実践

two sum (hash で O(n))

このレッスンで分かること

「配列 nums から、和が target になる 2 つの異なるインデックス を見つけたい」というのが two sum 問題です

関数 twoSum(nums, target) は次を満たします

「答えが存在しないケース」は今回は考慮しません

two sum (hash で O(n)) とは

整数配列と目標値 target を受け取り、和が target になる 2 要素のインデックスをソート済みリストで返す。

「配列 nums から、和が target になる 2 つの異なるインデックス を見つけたい」というのが two sum 問題です。コーディング面接の超定番で、hashmap の威力を一番手っ取り早く実感できる問題でもあります。素朴に書けば二重ループで O(n^2)、hashmap を使えば 1 パスで O(n)。本レッスンではこの差を体感してください。

two sum は hashmap を覚えるための最良の問題。1 パスで O(n) に落とすパターンを身につける。

仕様

関数 twoSum(nums, target) は次を満たします。

入力は整数配列 nums と目標値 target
nums[i] + nums[j] == target かつ i != j を満たす [i, j] を返す
必ずちょうど 1 組の答えが存在する
戻り値は 昇順にソートしたインデックスペア (例: [0, 3])

「答えが存在しないケース」は今回は考慮しません。テストケースもすべて 1 つの解を持つ前提です。

素朴な O(n^2) の解

Python

def twoSum(nums, target):
    n = len(nums)
    for i in range(n):
        for j in range(i + 1, n):
            if nums[i] + nums[j] == target:
                return [i, j]
    return []

この実装は明快ですが、要素が 10 万件あると 10^10 回ループするので現実的に終わりません。

O(n) に落とすアイデア

各要素 nums[i] を見たとき、「target - nums[i] という値が既に登場していれば、それと組めば答え」と考えます。これを hashmap で記録すれば、過去の出現位置を O(1) で引けます。

diagram (will load when visible)

Python での実装

Python

def twoSum(nums, target):
    seen = {}
    for i, v in enumerate(nums):
        need = target - v
        if need in seen:
            a, b = seen[need], i
            return sorted([a, b])
        seen[v] = i
    return []

enumerate で i, v を取り、need が既に seen に入っていれば即座に答え。なければ自分自身を seen に登録します。自分自身を先に登録すると同一インデックスを 2 回使うバグが入る ので、必ず「探してから登録」の順番にします。

JavaScript での実装

JavaScript

function twoSum(nums, target) {
  const seen = new Map();
  for (let i = 0; i < nums.length; i++) {
    const need = target - nums[i];
    if (seen.has(need)) {
      const a = seen.get(need);
      return [a, i].sort((x, y) => x - y);
    }
    seen.set(nums[i], i);
  }
  return [];
}

JS でも基本構造は同じです。Map を使うと整数キーをそのまま扱えるので、Object よりこちらが安心です。

Java での実装

Java

import java.util.*;

public class Solution {
    public static int[] twoSum(int[] nums, int target) {
        Map<Integer, Integer> seen = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            int need = target - nums[i];
            if (seen.containsKey(need)) {
                int a = seen.get(need);
                int[] res = { a, i };
                Arrays.sort(res);
                return res;
            }
            seen.put(nums[i], i);
        }
        return new int[0];
    }
}

よくある間違い

1 つ目は 登録順 です。seen[v] = i を if need in seen より先に書くと、nums = [3, 3], target = 6 のようなケースで i == j の答えを返してしまいます。
2 つ目は 戻り値の順序。[i, j] のうち、出会った順では j > i ですが、テスト側で昇順を期待していると逆順だと落ちます。明示的に sort するのが安全です。
3 つ目は インデックスではなく値を返してしまう こと。問題文をよく読んで、何を返すか確認すること。

集合系の問題でインデックスを返すときは、必ず仕様で「順序」が決まっているかを確認する。

やってみよう

nums = [2, 7, 11, 15], target = 9 で [0, 1] が返ることを確認する。
target を満たすペアが複数あるバージョンを考えて、最初に見つかった 1 つを返す実装に書き換えてみる。
同じ問題を 3 sum (3 要素の和) に拡張する。基本は 1 つ固定して残りを 2 sum で解く。

two sum を O(n) で書けるようになれば、hashmap の使いどころが感覚として身につく。

よくある質問

Q. なぜ hashmap で O(n) になるのですか？

A. 1 要素を見る間に「相棒の値が既に出たか」を hash の O(1) ルックアップで判定できるからです。二重ループ O(n²) を 1 周 + O(1) ルックアップに置き換えるのが two sum 系の発想です。空間 O(n) を使う代わりに時間を大幅短縮します。

Q. ハッシュ衝突が起きるとどうなりますか？

A. 同じバケットに複数要素が入りリスト走査になります。ロードファクタが高くなるとリサイズが走り、内部的に再ハッシュが行われます。最悪 O(n) になりますが、現代の hashmap は耐衝突ハッシュ関数で実用的にはほぼ O(1) です。

Q. キーに使える型は何ですか？

A. イミュータブル（変更不可）でハッシュ可能な型です。文字列・数値・タプルは OK ですが、リストや辞書は NG です。自作クラスは equals/hashCode（Java）または hash/eq（Python）を実装する必要があります。

次のレッスン

次は 部分配列の和 = k の個数 で、整数配列と k を受け取り、和が k になる連続する部分配列の個数を hashmap と累積和で O(n) で数える手法を学びます。

事前確認 — 進む前に次の 3 つができることを確認しましょう。

twoSum の要点を自分の言葉で説明できる
このレッスンの最小コード (または操作手順) を見ずに書ける
練習問題やクイズで間違えた箇所を読み直して理解した

理解度チェック (30 秒)

Q. twoSum とは何か、1 文で説明してください。

この章のポイント

A. 本文の「このレッスンで分かること」または冒頭の説明文を見直し、自分の言葉で要約できれば OK。詰まったら本レッスンの最初の H2 セクションを読み返してみましょう。

関連レッスン

two sum (hash で O(n))