Swift教室 Swift と競技プログラミングの教室

11. Two Sum

読了目安 約1分

HashMap / LeetCode 1 (Easy) — 値から添字を引く辞書で、和が target になる 2 数を 1 回の走査で探す。

この章の目次

問題

LeetCode 1. Two Sum(Easy / カテゴリ: HashMap)

整数の配列 nums と整数 target が与えられます。 和がちょうど target になる 2 つの要素を選び、その添字の組を返します。 答えは必ず 1 組だけ存在し、同じ要素を 2 回使うことはできません。

考え方

ヒント 1

2 要素の組を全部試すと O(n²) です。 「いま見ている数の相方が、もう登場したか」を一発で調べられれば、走査 1 回で済みます。

ヒント 2

値から添字を引ける辞書を、走査しながら育てます。 nums[i] を見たとき、辞書に target - nums[i] があれば、その添字と i が答えです。 なければ nums[i]i を辞書に登録して次へ進みます。

Swift 実装のポイント

  • 辞書は [Int: Int] で「値 → 添字」を持ちます。
  • if let j = seen[target - num] と書くと、相方の存在チェックと添字の取り出しが同時にできます。
  • 添字と値の同時走査は for (i, num) in nums.enumerated() が定番です。
模範解答
Swift
class Solution {
    func twoSum(_ nums: [Int], _ target: Int) -> [Int] {
        var seen: [Int: Int] = [:]  // 値 → 添字
        for (i, num) in nums.enumerated() {
            if let j = seen[target - num] {
                return [j, i]
            }
            seen[num] = i
        }
        return []
    }
}

計算量: O(n)(辞書の検索と挿入は平均 O(1) で、配列を 1 回走査するだけです)。