LeetCode - Search in Rotated Sorted ArrayをSwiftで解説する


  • 元の問題は英語なので、代わりに翻訳したものを記載しています。
  • 回答はあくまでも例であり、他の方法も存在します。

前置き

少し難易度はありますが、Swift で binary search を使ったことがない人にとっては、勉強するのに良い問題でしょう。

問題

問題 難易度
Search in Rotated Sorted Array medium
昇順の整数配列'nums'と整数'target'が与えられます。

'nums'があらかじめ知らないピボットで回転されているとします。
(例えば、[0,1,2,4,5,6,7] は [4,5,6,7,0,1,2] になるかもしれません)

'target'が配列の中で見つかった場合は、そのインデックスを返し、そうでない場合は -1 を返します。

つまり、以下の工程を行うことで結果を得ることができます。

① 配列(nums)から指定された数字(target)を探す
② 見つけた数字のインデックス番号を返す
③ 見つからなければ-1を返す

答え

  • Swift のAPIを使用する方法
  • 二分探索(binary search)を用いる方法

の2つを紹介します。

1. Swift のAPIを使用する方法

Swift には便利な firstIndex(where: /* 条件 */ } という関数が用意されており、こちらを使用するとワンライナーで書くことができます。

func search(_ nums: [Int], _ target: Int) -> Int {
    return nums.firstIndex(where: {$0 == target}) ?? -1
}

この関数は、探している整数(target)が見つからなければ nil を返すようになっています。
なので ?? -1-1 を返すようにしています。

ただし、こちらを使用するのは競技プログラミングにおいては反則のようなものなので、アルゴリズムを使ったものを見ていきましょう。
(用意されたAPIを使用して解いてしまっては、アルゴリズムの勉強にならないので)

2. 二分探索(binary search)を用いる方法

二分探索というのは、普通に探索すると時間がかかってしまうものを、範囲を絞って探索することで早く見つけられるアルゴリズムになります。

今回は詳しく説明しないので、以下を見ておくと良いでしょう。

二分探索を適用したコードは以下のようになります。

func search(_ nums: [Int], _ target: Int) -> Int {
    var left = 0
    var right = nums.count - 1
    
    while left <= right {
        let mid = (left+right)/2
        
        /// 値が見つかった場合は結果を返す
        guard nums[mid] != target else { return mid }
        
        /// 左が単調増加
        if nums[left] <= nums[mid] {
            /// targetが範囲内にあるかどうか
            if nums[left] <= target && target < nums[mid] {
                right = mid-1
            } else {
                left = mid+1
            }
         
        /// 右が単調増加
        } else {
            /// targetが範囲内にあるかどうか
            if nums[mid] < target && target <= nums[right] {
                left = mid+1
            } else {
                right = mid-1
            }
        }
    }
    
    /// 該当しない場合
    return -1
}

今回の問題では
「昇順の整数配列があらかじめピボットで回転されている」
という条件があるので、こちらを考慮することが求められています。

では、コードを細かく見ていきます。

① 二分探索の大枠部分

解答はいくつか存在しますが、核となる二分探索のコード部分はだいたい同じような書き方になるでしょう。

var left = 0
var right = nums.count - 1

while left <= right {
    let mid = (left+right)/2
    
    if nums[left] <= nums[mid] {
        // update left or right
    } else {
        // update left or right
    }
}

与えられた target を探すために、条件の「昇順である」ということを利用します。なぜなら、昇順の配列を見つけることで、 target がその配列内にあるかどうかを探すことができるからです。
filter 関数などは使わない前提で)

つまり、二分探索で間をとった場合に、どちらが単調増加(昇順を保ったまま)なのかを判断していくことになります。

② 範囲を狭めていく部分

昇順の配列が判断できた場合、配列内に target があるのかを見ていきます。

  • が昇順の配列の場合
/// 配列内にtargetが存在するかどうか
if nums[left] <= target && target < nums[mid] {
    /// 存在する場合は、右側を中間点までの範囲に狭める
    right = mid-1
} else {
    /// 存在しない場合、反対側の配列に存在する可能性があるので、左側を中間点までの範囲に狭める
    left = mid+1
}
  • が昇順の配列の場合(左の逆番なのでコメントは書きません)
if nums[mid] < target && target <= nums[right] {
    left = mid+1
} else {
    right = mid-1
}

③ 値を発見する部分

範囲を狭めていき、中間点と値が一致したものを返します。

guard nums[mid] != target else { return mid }

1点、guard で記載してしまったので少しわかりずらかったでしょう。

if 文で書き直すと

if nums[mid] == target {
    return mid
}

となり、少しわかりやすくなったでしょう。(好みの問題ですが)

また、最後に該当するものがない場合を考慮する必要があり、

return -1

という形で結果を返してあげます。

実行結果

Kind Runtime Memory
① firstIndex(where: { }) 16ms~ 14MB~
② binary search 12ms~ 14MB~

引用