modalsoul’s blog

これは“失敗”と呼べるかもしれないが、ぼくは“学習体験”と呼びたい

LeetCode 1. Two Sum

某入社エントリーに影響を受けてLeetCodeをはじめてみた

leetcode.com

最近プライベートで書くコードがスマートホーム系のちょっとしたPythonスクリプトばかりで飽きてきたので、Scalaの手直しも兼ねてやってみる

1. Two Sum

https://leetcode.com/problems/two-sum/

Intの配列が与えられるので、和が指定される値と等しくなる2つの数値のインデックスを返す

入力に対して必ず解が存在し、同じ値を2度使うことはない

Example:

Given nums = [2, 7, 11, 15], target = 9

Because nums[0] + nums[1] = 2 + 7 = 9,
return [0, 1].

考えたこと

  • リストなのでheadtailに分けて考えよう
  • headの値とtailのいずれかの値の和が、targetと等しくなった場合、headのインデックスとtailの値のインデックスを返せばOK
  • tailのいずれの値でもtargetと等しくならない場合、tailをさらにheadtailに分けて考える(以下繰り返し

コードその1

object Solution {
    def twoSum(nums: Array[Int], target: Int): Array[Int] = {
        def index(n1:Int, arr:Array[Int]):Option[Int] = {
          val res = arr.indexWhere(_ + n1 == target) + 1
          if(res != 0) Some(res) else None
        }

        def run(nums:Array[Int], offset:Int = 0):Array[Int] = {
          index(nums.head, nums.tail) match {
            case Some(n2) => Array(offset, n2 + offset)
            case _ => run(nums.tail, offset + 1)
          }
        }
        run(nums)
    }
}

解説

def index(n1:Int, arr:Array[Int]):Option[Int]

任意のn1との和がtargetと等しくなる値のインデックスのオプションを返す

Seq.indexWhereの戻り型をIntからOption[Int]へ変換するためのもの

Seq.indexWhereOption:Option[Int]のようなメソッドがあれば良かったんだが。。

def run(nums:Array[Int], offset:Int = 0):Array[Int]

numsのheadに対して条件を満たす値がtailに存在するか否か

indexの結果をパターンマッチして、Someの場合に解発見、Noneの場合再帰する

tailを遡るたびにindexが進むので、その回数をoffsetで持ち回る

考察

一応できたけれど、O(n2)

改善の余地あり

コードその2

object Solution {
    def twoSum(nums: Array[Int], target: Int): Array[Int] = {
        val map = nums.zipWithIndex.toMap
        nums.zipWithIndex.foldLeft(Array.empty[Int]){(z, n) =>
          if(!z.isEmpty) z
          else {
            val t = target - n._1
            map.get(t).fold(z){ x => if(x != n._2) Array(n._2, x) else z}
          }
        }
    }
}

解説

  • n番目の値のインデックスとn番目の値をtargetから引いた値のインデックスが解になる
  • 値とインデックスのマップを生成

これでO(n)

コードその3

object Solution {
    def twoSum(nums: Array[Int], target: Int): Array[Int] = {
        nums.zipWithIndex.foldRight((Map.empty[Int, Int], Array.empty[Int])) { (n, z) =>
          if(!z._2.isEmpty) z
          else {
            val (map, acc) = z
            val (num, index) = n
            (map + n, map.get(target - num).fold(acc){ x => if(x != index) Array(index, x) else acc})
          }
        }._2
    }
}

解説

  • 基本的な考え方はコードその2と同じ
  • 値とインデックスのマップ生成時に、和がtargetになる値の有無をチェックする