Zenn
🐢

LeetCode - Squares of a Sorted Array をSwiftで解説する

2020/12/15に公開
iOS
Swift
競技プログラミング
AtCoder
LeetCode
tech
  • 元の問題は英語なので、代わりに翻訳したものを記載しています。
  • 回答はあくまでも例であり、他の方法も存在します。

前置き

Swiftで Sort アルゴリズム を書いたことがない人にとっては、勉強するのに良い問題でしょう。

問題

問題 難易度
Squares of a Sorted Array easy 
整数の配列 nums が小さい順にソートされています。
各数値を2乗した配列を小さい順にソートして返します。

補足として、以下のような true を返すテストケースが与えられます。

Input: nums = [-4,-1,0,3,10]
Output: [0,1,9,16,100]

テストケースの計算過程はこうなります。

  1. 元の値　　 → [-4, -1, 0, 3, 10]
  2. ２乗する 　→ [16, 1, 0, 9, 100]
  3. ソートする → [0, 1, 9, 16, 100]

ポイントとしては、配列にマイナスの値がある場合、2乗した際に値がプラスに反転します。その際、最初は小さい順にソートされていた配列の順番が変わってしまうため、再度ソートする必要があります。

答え

  • Swift のAPIを使用する方法
  • Sort アルゴリズムを用いる方法

の２つを紹介します。

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

func sortedSquares(_ A: [Int]) -> [Int] {
    return A.map { $0 * $0 }.sorted()
}

Swiftでは sorted メソッドが用意されています。

sorted メソッドは、内部的には複数の Sort アルゴリズムを組み合わせています。そのため自前で Sort のアルゴリズムを実装するよりも、こちらを使用する方が高速でしょう。
（参考： Swiftではどのようにソートを行っているか）

しかし、今回はあえてアルゴリズムを使った方法を見ていきます。

2. Sort アルゴリズムを用いる方法

ソートのアルゴリズムはいくつか存在しますが、今回は２種類を紹介します。

① Bubble Sort（バブルソート）

隣り合う要素の値を比較し、条件に応じて交換するアルゴリズムです。

bubble sort
(引用：バブルソートでシンプルなベンチマーク)

では、実際のコードを見ていきましょう。

func sortedSquares(_ A: [Int]) -> [Int] {
    // ① ２乗した配列を作る
    var result = [Int]()
    for n in A {
        result.append(n*n)
    }
    
    // ② ソートしていく
    var isSwapped = true
    while isSwapped {
        // 一度でも値の交換をした場合は、再度比較をし直すためのフラグ
        isSwapped = false
        
        // 最初から数字を比較していく
        // 中の計算で'i-1'をしても問題ないように'０'ではなく'1'から始まっている
        for i in 1..<result.count {
            // 隣り合う数字の比較
            guard result[i] < result[i-1] else { continue }
            // 該当する場合は'swapAt'で値を交換する
            result.swapAt(i, i-1)
            isSwapped = true
        }
    }
    return result
}

再起関数使う書き方もありますが、今回はwhileとフラグを使って書いています。

今回こちらを紹介しましたが、Bubble Sort は総当たりで比較していくため、計算量が非常に多く時間がかかります。この問題で使用すると、残念ながらタイムアウトになってしまいます。

では、こちらよりも高速な Sort アルゴリズムを見てみましょう。

② Merge Sort（マージソート）

ソートされていないリストを2つのリストに分割し、それぞれをソートさせた後、それらを１つにする（マージする）アルゴリズムです。

bubble sort
(引用：ソースコード探検隊 » アルゴリズムとデータ構造 » アルゴリズム » マージソート)

では、実際のコードを見ていきましょう。

func sortedSquares(_ A: [Int]) -> [Int] {
    // ① ２乗した配列を作る
    var squaredNums = [Int]()
    for n in A {
        squaredNums.append(n*n)
    }
    
    // ② ソートしていく　&　③ 合成する 
    return mergeSort(squaredNums)
}

func mergeSort(_ lists: [Int]) -> [Int] {
    guard 1 < lists.count else { return lists }
    
    // ②-1 配列を２つに分割
    let count = lists.count
    let left = Array(lists[..<Int(count/2)])
    
    // ②-2 分割した個々の配列をソートし続ける
    let right = Array(lists[count/2...count-1])
    return merge(left: mergeSort(left), right: mergeSort(right))
}

func merge(left: [Int], right: [Int]) -> [Int] {
    var left = left, right = right, result = [Int]()
    
    // ③-1 ２つの配列から最初の値を見て、小さい値から順にリストに追加していくことで、２つの配列を合成していく
    while 0 < left.count && 0 < right.count {
        // `removeFirst`で追加した方の先頭の値を除外していく
        if left[0] <= right[0] {
            result.append(left.removeFirst())
        } else {
            result.append(right.removeFirst())
        }
    }
    
    // ③-2 `while`ではいずれかが'０'になった時に処理が終わってしまうので、最後の値を配列に追加する
    result.append(contentsOf: left)
    result.append(contentsOf: right)
    
    return result
}

Merge Sort では「 sort する関数」と「 merge する関数」をそれぞれ用意します。一見するとコード量は多いですが、Bubble Sort とは違い、個々の配列を見ることで比較する値の計算量が減少されます。

③ Quick Sort（クイックソート）

ソートされていないリストを、ピボットと呼ぶ要素を軸に分割を繰り返して整列を行うアルゴリズムです。

また、「分割を繰り返して整列を行う」ような手法を分割統治法 divide-and-conquer と呼びます。

quick-sort
(引用：ソースコード探検隊 » アルゴリズムとデータ構造 » アルゴリズム » クイックソート)

これはマージソートに似ており、マージソートも分割統治法の１つです。ただ、分割方法に違いがあります。

では、実際のコードを見ていきましょう。

func sortedSquares(_ A: [Int]) -> [Int] {
    // ① ２乗した配列を作る
    var squaredNums = [Int]()
    for n in A {
        squaredNums.append(n*n)
    }
    
    // ② ソートしていく 
    return quickSort(squaredNums)
}

func quickSort(_ nums: [Int]) -> [Int] {
    // 数が２つ以上ない場合、ソートできないのでそのまま返す
    guard 1 < nums.count else { return nums }

    // 基準値を決める
    let pivot = nums[nums.count/2]
    // 基準値より小さいものを取り出す
    let less = nums.filter { $0 < pivot }
    // 基準値と同じものを取り出す
    let equal = nums.filter { $0 == pivot }
    // 基準値より大きいものを取り出す
    let greater = nums.filter { pivot < $0 }

    // 同じ処理を繰り返していく
    return quickSort(less) + equal + quickSort(greater)
}

再帰関数を使うことですごく短いコードになっています。ただし、「実行結果」にも記載してありますが、関数を何度も呼び出すため、メモリを多く使用します。

終わりに

3つの Sort を紹介しましたが、他にもアルゴリズムはいくつかあるので、興味がある方は調べてみてください。

実行結果

Kind Runtime Memory
sorted() 380ms~ 16MB前後
Bubble Sort Time Limit Exceeded -
Merge Sort 740ms~ 16MB前後
Quick Sort 380ms~ 75MB前後

※実行回数を複数回行うことで値は多少前後します

引用

Discussion